Search in sources :

Example 31 with List

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

the class RemoteExecutionTest method testMetrics.

@Test
@SuppressWarnings("FutureReturnValueIgnored")
public void testMetrics() throws Exception {
    launchSdkHarness(PipelineOptionsFactory.create());
    MetricsDoFn metricsDoFn = new MetricsDoFn();
    Pipeline p = Pipeline.create();
    PCollection<String> input = p.apply("impulse", Impulse.create()).apply("create", ParDo.of(metricsDoFn)).setCoder(StringUtf8Coder.of());
    SingleOutput<String, String> pardo = ParDo.of(new DoFn<String, String>() {

        @ProcessElement
        public void process(ProcessContext ctxt) {
            // Output the element twice to keep unique numbers in asserts, 6 output elements.
            ctxt.output(ctxt.element());
            ctxt.output(ctxt.element());
        }
    });
    input.apply("processA", pardo).setCoder(StringUtf8Coder.of());
    input.apply("processB", pardo).setCoder(StringUtf8Coder.of());
    RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p);
    FusedPipeline fused = GreedyPipelineFuser.fuse(pipelineProto);
    Optional<ExecutableStage> optionalStage = Iterables.tryFind(fused.getFusedStages(), (ExecutableStage stage) -> true);
    checkState(optionalStage.isPresent(), "Expected a stage with side inputs.");
    ExecutableStage stage = optionalStage.get();
    ExecutableProcessBundleDescriptor descriptor = ProcessBundleDescriptors.fromExecutableStage("test_stage", stage, dataServer.getApiServiceDescriptor(), stateServer.getApiServiceDescriptor());
    BundleProcessor processor = controlClient.getProcessor(descriptor.getProcessBundleDescriptor(), descriptor.getRemoteInputDestinations(), stateDelegator);
    Map<String, Coder> remoteOutputCoders = descriptor.getRemoteOutputCoders();
    Map<String, RemoteOutputReceiver<?>> outputReceivers = new HashMap<>();
    for (Entry<String, Coder> remoteOutputCoder : remoteOutputCoders.entrySet()) {
        List<WindowedValue<?>> outputContents = Collections.synchronizedList(new ArrayList<>());
        outputReceivers.put(remoteOutputCoder.getKey(), RemoteOutputReceiver.of((Coder<WindowedValue<?>>) remoteOutputCoder.getValue(), outputContents::add));
    }
    final String testPTransformId = "create-ParMultiDo-Metrics-";
    BundleProgressHandler progressHandler = new BundleProgressHandler() {

        @Override
        public void onProgress(ProcessBundleProgressResponse response) {
            MetricsDoFn.ALLOW_COMPLETION.get(metricsDoFn.uuid).countDown();
            List<Matcher<MonitoringInfo>> matchers = new ArrayList<>();
            // We expect all user counters except for the ones in @FinishBundle
            // Since non-user metrics are registered at bundle creation time, they will still report
            // values most of which will be 0.
            SimpleMonitoringInfoBuilder builder = new SimpleMonitoringInfoBuilder();
            builder.setUrn(MonitoringInfoConstants.Urns.USER_SUM_INT64).setLabel(MonitoringInfoConstants.Labels.NAMESPACE, RemoteExecutionTest.class.getName()).setLabel(MonitoringInfoConstants.Labels.NAME, MetricsDoFn.PROCESS_USER_COUNTER_NAME);
            builder.setLabel(MonitoringInfoConstants.Labels.PTRANSFORM, testPTransformId);
            builder.setInt64SumValue(1);
            matchers.add(MonitoringInfoMatchers.matchSetFields(builder.build()));
            builder = new SimpleMonitoringInfoBuilder();
            builder.setUrn(MonitoringInfoConstants.Urns.USER_SUM_INT64).setLabel(MonitoringInfoConstants.Labels.NAMESPACE, RemoteExecutionTest.class.getName()).setLabel(MonitoringInfoConstants.Labels.NAME, MetricsDoFn.START_USER_COUNTER_NAME);
            builder.setLabel(MonitoringInfoConstants.Labels.PTRANSFORM, testPTransformId);
            builder.setInt64SumValue(10);
            matchers.add(MonitoringInfoMatchers.matchSetFields(builder.build()));
            builder = new SimpleMonitoringInfoBuilder();
            builder.setUrn(MonitoringInfoConstants.Urns.USER_SUM_INT64).setLabel(MonitoringInfoConstants.Labels.NAMESPACE, RemoteExecutionTest.class.getName()).setLabel(MonitoringInfoConstants.Labels.NAME, MetricsDoFn.FINISH_USER_COUNTER_NAME);
            builder.setLabel(MonitoringInfoConstants.Labels.PTRANSFORM, testPTransformId);
            matchers.add(not(MonitoringInfoMatchers.matchSetFields(builder.build())));
            // User Distributions.
            builder.setUrn(MonitoringInfoConstants.Urns.USER_DISTRIBUTION_INT64).setLabel(MonitoringInfoConstants.Labels.NAMESPACE, RemoteExecutionTest.class.getName()).setLabel(MonitoringInfoConstants.Labels.NAME, MetricsDoFn.PROCESS_USER_DISTRIBUTION_NAME);
            builder.setLabel(MonitoringInfoConstants.Labels.PTRANSFORM, testPTransformId);
            builder.setInt64DistributionValue(DistributionData.create(1, 1, 1, 1));
            matchers.add(MonitoringInfoMatchers.matchSetFields(builder.build()));
            builder = new SimpleMonitoringInfoBuilder();
            builder.setUrn(MonitoringInfoConstants.Urns.USER_DISTRIBUTION_INT64).setLabel(MonitoringInfoConstants.Labels.NAMESPACE, RemoteExecutionTest.class.getName()).setLabel(MonitoringInfoConstants.Labels.NAME, MetricsDoFn.START_USER_DISTRIBUTION_NAME);
            builder.setLabel(MonitoringInfoConstants.Labels.PTRANSFORM, testPTransformId);
            builder.setInt64DistributionValue(DistributionData.create(10, 1, 10, 10));
            matchers.add(MonitoringInfoMatchers.matchSetFields(builder.build()));
            builder = new SimpleMonitoringInfoBuilder();
            builder.setUrn(MonitoringInfoConstants.Urns.USER_DISTRIBUTION_INT64).setLabel(MonitoringInfoConstants.Labels.NAMESPACE, RemoteExecutionTest.class.getName()).setLabel(MonitoringInfoConstants.Labels.NAME, MetricsDoFn.FINISH_USER_DISTRIBUTION_NAME);
            builder.setLabel(MonitoringInfoConstants.Labels.PTRANSFORM, testPTransformId);
            matchers.add(not(MonitoringInfoMatchers.matchSetFields(builder.build())));
            assertThat(response.getMonitoringInfosList(), Matchers.hasItems(matchers.toArray(new Matcher[0])));
        }

        @Override
        public void onCompleted(ProcessBundleResponse response) {
            List<Matcher<MonitoringInfo>> matchers = new ArrayList<>();
            // User Counters.
            SimpleMonitoringInfoBuilder builder = new SimpleMonitoringInfoBuilder();
            builder.setUrn(MonitoringInfoConstants.Urns.USER_SUM_INT64).setLabel(MonitoringInfoConstants.Labels.NAMESPACE, RemoteExecutionTest.class.getName()).setLabel(MonitoringInfoConstants.Labels.NAME, MetricsDoFn.PROCESS_USER_COUNTER_NAME);
            builder.setLabel(MonitoringInfoConstants.Labels.PTRANSFORM, testPTransformId);
            builder.setInt64SumValue(1);
            matchers.add(MonitoringInfoMatchers.matchSetFields(builder.build()));
            builder = new SimpleMonitoringInfoBuilder();
            builder.setUrn(MonitoringInfoConstants.Urns.USER_SUM_INT64).setLabel(MonitoringInfoConstants.Labels.NAMESPACE, RemoteExecutionTest.class.getName()).setLabel(MonitoringInfoConstants.Labels.NAME, MetricsDoFn.START_USER_COUNTER_NAME);
            builder.setLabel(MonitoringInfoConstants.Labels.PTRANSFORM, testPTransformId);
            builder.setInt64SumValue(10);
            matchers.add(MonitoringInfoMatchers.matchSetFields(builder.build()));
            builder = new SimpleMonitoringInfoBuilder();
            builder.setUrn(MonitoringInfoConstants.Urns.USER_SUM_INT64).setLabel(MonitoringInfoConstants.Labels.NAMESPACE, RemoteExecutionTest.class.getName()).setLabel(MonitoringInfoConstants.Labels.NAME, MetricsDoFn.FINISH_USER_COUNTER_NAME);
            builder.setLabel(MonitoringInfoConstants.Labels.PTRANSFORM, testPTransformId);
            builder.setInt64SumValue(100);
            matchers.add(MonitoringInfoMatchers.matchSetFields(builder.build()));
            // User Distributions.
            builder.setUrn(MonitoringInfoConstants.Urns.USER_DISTRIBUTION_INT64).setLabel(MonitoringInfoConstants.Labels.NAMESPACE, RemoteExecutionTest.class.getName()).setLabel(MonitoringInfoConstants.Labels.NAME, MetricsDoFn.PROCESS_USER_DISTRIBUTION_NAME);
            builder.setLabel(MonitoringInfoConstants.Labels.PTRANSFORM, testPTransformId);
            builder.setInt64DistributionValue(DistributionData.create(1, 1, 1, 1));
            matchers.add(MonitoringInfoMatchers.matchSetFields(builder.build()));
            builder = new SimpleMonitoringInfoBuilder();
            builder.setUrn(MonitoringInfoConstants.Urns.USER_DISTRIBUTION_INT64).setLabel(MonitoringInfoConstants.Labels.NAMESPACE, RemoteExecutionTest.class.getName()).setLabel(MonitoringInfoConstants.Labels.NAME, MetricsDoFn.START_USER_DISTRIBUTION_NAME);
            builder.setLabel(MonitoringInfoConstants.Labels.PTRANSFORM, testPTransformId);
            builder.setInt64DistributionValue(DistributionData.create(10, 1, 10, 10));
            matchers.add(MonitoringInfoMatchers.matchSetFields(builder.build()));
            builder = new SimpleMonitoringInfoBuilder();
            builder.setUrn(MonitoringInfoConstants.Urns.USER_DISTRIBUTION_INT64).setLabel(MonitoringInfoConstants.Labels.NAMESPACE, RemoteExecutionTest.class.getName()).setLabel(MonitoringInfoConstants.Labels.NAME, MetricsDoFn.FINISH_USER_DISTRIBUTION_NAME);
            builder.setLabel(MonitoringInfoConstants.Labels.PTRANSFORM, testPTransformId);
            builder.setInt64DistributionValue(DistributionData.create(100, 1, 100, 100));
            matchers.add(MonitoringInfoMatchers.matchSetFields(builder.build()));
            // The element counter should be counted only once for the pcollection.
            // So there should be only two elements.
            builder = new SimpleMonitoringInfoBuilder();
            builder.setUrn(MonitoringInfoConstants.Urns.ELEMENT_COUNT);
            builder.setLabel(MonitoringInfoConstants.Labels.PCOLLECTION, "impulse.out");
            builder.setInt64SumValue(1);
            matchers.add(MonitoringInfoMatchers.matchSetFields(builder.build()));
            builder = new SimpleMonitoringInfoBuilder();
            builder.setUrn(MonitoringInfoConstants.Urns.ELEMENT_COUNT);
            builder.setLabel(MonitoringInfoConstants.Labels.PCOLLECTION, "create/ParMultiDo(Metrics).output");
            builder.setInt64SumValue(3);
            matchers.add(MonitoringInfoMatchers.matchSetFields(builder.build()));
            // Verify that the element count is not double counted if two PCollections consume it.
            builder = new SimpleMonitoringInfoBuilder();
            builder.setUrn(MonitoringInfoConstants.Urns.ELEMENT_COUNT);
            builder.setLabel(MonitoringInfoConstants.Labels.PCOLLECTION, "processA/ParMultiDo(Anonymous).output");
            builder.setInt64SumValue(6);
            matchers.add(MonitoringInfoMatchers.matchSetFields(builder.build()));
            builder = new SimpleMonitoringInfoBuilder();
            builder.setUrn(MonitoringInfoConstants.Urns.ELEMENT_COUNT);
            builder.setLabel(MonitoringInfoConstants.Labels.PCOLLECTION, "processB/ParMultiDo(Anonymous).output");
            builder.setInt64SumValue(6);
            matchers.add(MonitoringInfoMatchers.matchSetFields(builder.build()));
            // Check for execution time metrics for the testPTransformId
            builder = new SimpleMonitoringInfoBuilder();
            builder.setUrn(MonitoringInfoConstants.Urns.START_BUNDLE_MSECS);
            builder.setType(TypeUrns.SUM_INT64_TYPE);
            builder.setLabel(MonitoringInfoConstants.Labels.PTRANSFORM, testPTransformId);
            matchers.add(allOf(MonitoringInfoMatchers.matchSetFields(builder.build()), MonitoringInfoMatchers.counterValueGreaterThanOrEqualTo(1)));
            // Check for execution time metrics for the testPTransformId
            builder = new SimpleMonitoringInfoBuilder();
            builder.setUrn(Urns.PROCESS_BUNDLE_MSECS);
            builder.setType(TypeUrns.SUM_INT64_TYPE);
            builder.setLabel(MonitoringInfoConstants.Labels.PTRANSFORM, testPTransformId);
            matchers.add(allOf(MonitoringInfoMatchers.matchSetFields(builder.build()), MonitoringInfoMatchers.counterValueGreaterThanOrEqualTo(2)));
            builder = new SimpleMonitoringInfoBuilder();
            builder.setUrn(Urns.FINISH_BUNDLE_MSECS);
            builder.setType(TypeUrns.SUM_INT64_TYPE);
            builder.setLabel(MonitoringInfoConstants.Labels.PTRANSFORM, testPTransformId);
            matchers.add(allOf(MonitoringInfoMatchers.matchSetFields(builder.build()), MonitoringInfoMatchers.counterValueGreaterThanOrEqualTo(3)));
            assertThat(response.getMonitoringInfosList(), Matchers.hasItems(matchers.toArray(new Matcher[0])));
        }
    };
    ExecutorService executor = Executors.newSingleThreadExecutor();
    try (RemoteBundle bundle = processor.newBundle(outputReceivers, StateRequestHandler.unsupported(), progressHandler)) {
        Iterables.getOnlyElement(bundle.getInputReceivers().values()).accept(valueInGlobalWindow(CoderUtils.encodeToByteArray(StringUtf8Coder.of(), "X")));
        executor.submit(() -> {
            checkState(MetricsDoFn.AFTER_PROCESS.get(metricsDoFn.uuid).await(60, TimeUnit.SECONDS), "Runner waited too long for DoFn to get to AFTER_PROCESS.");
            bundle.requestProgress();
            return (Void) null;
        });
    }
    executor.shutdown();
}
Also used : ConcurrentHashMap(java.util.concurrent.ConcurrentHashMap) HashMap(java.util.HashMap) Matcher(org.hamcrest.Matcher) ArrayList(java.util.ArrayList) ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) RunnerApi(org.apache.beam.model.pipeline.v1.RunnerApi) SimpleMonitoringInfoBuilder(org.apache.beam.runners.core.metrics.SimpleMonitoringInfoBuilder) BundleProcessor(org.apache.beam.runners.fnexecution.control.SdkHarnessClient.BundleProcessor) WindowedValue(org.apache.beam.sdk.util.WindowedValue) ExecutableStage(org.apache.beam.runners.core.construction.graph.ExecutableStage) ExecutableProcessBundleDescriptor(org.apache.beam.runners.fnexecution.control.ProcessBundleDescriptors.ExecutableProcessBundleDescriptor) ProcessBundleResponse(org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleResponse) KvCoder(org.apache.beam.sdk.coders.KvCoder) Coder(org.apache.beam.sdk.coders.Coder) StringUtf8Coder(org.apache.beam.sdk.coders.StringUtf8Coder) BigEndianLongCoder(org.apache.beam.sdk.coders.BigEndianLongCoder) FusedPipeline(org.apache.beam.runners.core.construction.graph.FusedPipeline) Pipeline(org.apache.beam.sdk.Pipeline) FusedPipeline(org.apache.beam.runners.core.construction.graph.FusedPipeline) ScheduledExecutorService(java.util.concurrent.ScheduledExecutorService) ExecutorService(java.util.concurrent.ExecutorService) ProcessBundleProgressResponse(org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleProgressResponse) Test(org.junit.Test)

Example 32 with List

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

the class RemoteExecutionTest method testSplit.

@Test(timeout = 60000L)
public void testSplit() throws Exception {
    launchSdkHarness(PipelineOptionsFactory.create());
    Pipeline p = Pipeline.create();
    p.apply("impulse", Impulse.create()).apply("create", ParDo.of(new DoFn<byte[], String>() {

        @ProcessElement
        public void process(ProcessContext ctxt) {
            ctxt.output("zero");
            ctxt.output(WaitingTillSplitRestrictionTracker.WAIT_TILL_SPLIT);
            ctxt.output("two");
        }
    })).apply("forceSplit", ParDo.of(new DoFn<String, String>() {

        @GetInitialRestriction
        public String getInitialRestriction(@Element String element) {
            return element;
        }

        @NewTracker
        public WaitingTillSplitRestrictionTracker newTracker(@Restriction String restriction) {
            return new WaitingTillSplitRestrictionTracker(restriction);
        }

        @ProcessElement
        public void process(RestrictionTracker<String, Void> tracker, ProcessContext context) {
            while (tracker.tryClaim(null)) {
            }
            context.output(tracker.currentRestriction());
        }
    })).apply("addKeys", WithKeys.of("foo")).setCoder(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())).apply("gbk", GroupByKey.create());
    RunnerApi.Pipeline pipeline = PipelineTranslation.toProto(p);
    // Expand any splittable DoFns within the graph to enable sizing and splitting of bundles.
    RunnerApi.Pipeline pipelineWithSdfExpanded = ProtoOverrides.updateTransform(PTransformTranslation.PAR_DO_TRANSFORM_URN, pipeline, SplittableParDoExpander.createSizedReplacement());
    FusedPipeline fused = GreedyPipelineFuser.fuse(pipelineWithSdfExpanded);
    // Find the fused stage with the SDF ProcessSizedElementAndRestriction transform
    Optional<ExecutableStage> optionalStage = Iterables.tryFind(fused.getFusedStages(), (ExecutableStage stage) -> Iterables.filter(stage.getTransforms(), (PTransformNode node) -> PTransformTranslation.SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN.equals(node.getTransform().getSpec().getUrn())).iterator().hasNext());
    checkState(optionalStage.isPresent(), "Expected a stage with SDF ProcessSizedElementAndRestriction.");
    ExecutableStage stage = optionalStage.get();
    ExecutableProcessBundleDescriptor descriptor = ProcessBundleDescriptors.fromExecutableStage("my_stage", stage, dataServer.getApiServiceDescriptor());
    BundleProcessor processor = controlClient.getProcessor(descriptor.getProcessBundleDescriptor(), descriptor.getRemoteInputDestinations());
    Map<String, ? super Coder<WindowedValue<?>>> remoteOutputCoders = descriptor.getRemoteOutputCoders();
    Map<String, Collection<? super WindowedValue<?>>> outputValues = new HashMap<>();
    Map<String, RemoteOutputReceiver<?>> outputReceivers = new HashMap<>();
    for (Entry<String, ? super Coder<WindowedValue<?>>> remoteOutputCoder : remoteOutputCoders.entrySet()) {
        List<? super WindowedValue<?>> outputContents = Collections.synchronizedList(new ArrayList<>());
        outputValues.put(remoteOutputCoder.getKey(), outputContents);
        outputReceivers.put(remoteOutputCoder.getKey(), RemoteOutputReceiver.of((Coder) remoteOutputCoder.getValue(), (FnDataReceiver<? super WindowedValue<?>>) outputContents::add));
    }
    List<ProcessBundleSplitResponse> splitResponses = new ArrayList<>();
    List<ProcessBundleResponse> checkpointResponses = new ArrayList<>();
    List<String> requestsFinalization = new ArrayList<>();
    ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();
    ScheduledFuture<Object> future;
    // Execute the remote bundle.
    try (RemoteBundle bundle = processor.newBundle(outputReceivers, Collections.emptyMap(), StateRequestHandler.unsupported(), BundleProgressHandler.ignored(), splitResponses::add, checkpointResponses::add, requestsFinalization::add)) {
        Iterables.getOnlyElement(bundle.getInputReceivers().values()).accept(valueInGlobalWindow(sdfSizedElementAndRestrictionForTest(WaitingTillSplitRestrictionTracker.WAIT_TILL_SPLIT)));
        // Keep sending splits until the bundle terminates.
        future = (ScheduledFuture) executor.scheduleWithFixedDelay(() -> bundle.split(0.5), 0L, 100L, TimeUnit.MILLISECONDS);
    }
    future.cancel(false);
    executor.shutdown();
    assertTrue(requestsFinalization.isEmpty());
    assertTrue(checkpointResponses.isEmpty());
    // We only validate the last split response since it is the only one that could possibly
    // contain the SDF split, all others will be a reduction in the ChannelSplit range.
    assertFalse(splitResponses.isEmpty());
    ProcessBundleSplitResponse splitResponse = splitResponses.get(splitResponses.size() - 1);
    ChannelSplit channelSplit = Iterables.getOnlyElement(splitResponse.getChannelSplitsList());
    // There is only one outcome for the final split that can happen since the SDF is blocking the
    // bundle from completing and hence needed to be split.
    assertEquals(-1L, channelSplit.getLastPrimaryElement());
    assertEquals(1L, channelSplit.getFirstResidualElement());
    assertEquals(1, splitResponse.getPrimaryRootsCount());
    assertEquals(1, splitResponse.getResidualRootsCount());
    assertThat(Iterables.getOnlyElement(outputValues.values()), containsInAnyOrder(valueInGlobalWindow(KV.of("foo", WaitingTillSplitRestrictionTracker.PRIMARY))));
}
Also used : ConcurrentHashMap(java.util.concurrent.ConcurrentHashMap) HashMap(java.util.HashMap) PTransformNode(org.apache.beam.runners.core.construction.graph.PipelineNode.PTransformNode) ArrayList(java.util.ArrayList) ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) RunnerApi(org.apache.beam.model.pipeline.v1.RunnerApi) BundleProcessor(org.apache.beam.runners.fnexecution.control.SdkHarnessClient.BundleProcessor) WindowedValue(org.apache.beam.sdk.util.WindowedValue) ExecutableStage(org.apache.beam.runners.core.construction.graph.ExecutableStage) ExecutableProcessBundleDescriptor(org.apache.beam.runners.fnexecution.control.ProcessBundleDescriptors.ExecutableProcessBundleDescriptor) ProcessBundleResponse(org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleResponse) KvCoder(org.apache.beam.sdk.coders.KvCoder) Coder(org.apache.beam.sdk.coders.Coder) StringUtf8Coder(org.apache.beam.sdk.coders.StringUtf8Coder) BigEndianLongCoder(org.apache.beam.sdk.coders.BigEndianLongCoder) ScheduledExecutorService(java.util.concurrent.ScheduledExecutorService) FnDataReceiver(org.apache.beam.sdk.fn.data.FnDataReceiver) FusedPipeline(org.apache.beam.runners.core.construction.graph.FusedPipeline) Pipeline(org.apache.beam.sdk.Pipeline) FusedPipeline(org.apache.beam.runners.core.construction.graph.FusedPipeline) ProcessBundleSplitResponse(org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitResponse) DoFn(org.apache.beam.sdk.transforms.DoFn) Collection(java.util.Collection) PCollection(org.apache.beam.sdk.values.PCollection) ChannelSplit(org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitResponse.ChannelSplit) Test(org.junit.Test)

Example 33 with List

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.LogEntry.List 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));
}
Also used : FullWindowedValueCoder(org.apache.beam.sdk.util.WindowedValue.FullWindowedValueCoder) RemoteInputDestination(org.apache.beam.runners.fnexecution.data.RemoteInputDestination) ProcessBundleDescriptor(org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleDescriptor) BundleProcessor(org.apache.beam.runners.fnexecution.control.SdkHarnessClient.BundleProcessor) Test(org.junit.Test)

Example 34 with List

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

the class CommonCoderTest method instantiateCoder.

private static Coder<?> instantiateCoder(CommonCoder coder) {
    List<Coder<?>> components = new ArrayList<>();
    for (CommonCoder innerCoder : coder.getComponents()) {
        components.add(instantiateCoder(innerCoder));
    }
    // translator since we need to interact with a fake state client.
    if (coder.getUrn().equals(getUrn(StandardCoders.Enum.STATE_BACKED_ITERABLE))) {
        BeamFnStateClient stateClient = new BeamFnStateClient() {

            @Override
            public CompletableFuture<StateResponse> handle(StateRequest.Builder requestBuilder) {
                checkState(requestBuilder.hasGet());
                checkState(requestBuilder.hasStateKey());
                checkState(requestBuilder.getStateKey().hasRunner());
                StateResponse.Builder rval = StateResponse.newBuilder();
                rval.setId(requestBuilder.getId());
                rval.setGet(StateGetResponse.newBuilder().setData(coder.getState().getOrDefault(requestBuilder.getStateKey().getRunner().getKey(), ByteString.EMPTY)));
                return CompletableFuture.completedFuture(rval.build());
            }
        };
        return new StateBackedIterable.Coder<>(() -> Caches.noop(), stateClient, () -> "instructionId", (Coder) Iterables.getOnlyElement(components));
    }
    Class<? extends Coder> coderType = ImmutableBiMap.copyOf(new ModelCoderRegistrar().getCoderURNs()).inverse().get(coder.getUrn());
    checkNotNull(coderType, "Unknown coder URN: " + coder.getUrn());
    CoderTranslator<?> translator = new ModelCoderRegistrar().getCoderTranslators().get(coderType);
    checkNotNull(translator, "No translator found for common coder class: " + coderType.getSimpleName());
    return translator.fromComponents(components, coder.getPayload(), new TranslationContext() {
    });
}
Also used : DoubleCoder(org.apache.beam.sdk.coders.DoubleCoder) IntervalWindowCoder(org.apache.beam.sdk.transforms.windowing.IntervalWindow.IntervalWindowCoder) ByteCoder(org.apache.beam.sdk.coders.ByteCoder) KvCoder(org.apache.beam.sdk.coders.KvCoder) VarLongCoder(org.apache.beam.sdk.coders.VarLongCoder) BooleanCoder(org.apache.beam.sdk.coders.BooleanCoder) TimestampPrefixingWindowCoder(org.apache.beam.sdk.coders.TimestampPrefixingWindowCoder) Coder(org.apache.beam.sdk.coders.Coder) RowCoder(org.apache.beam.sdk.coders.RowCoder) StringUtf8Coder(org.apache.beam.sdk.coders.StringUtf8Coder) IterableCoder(org.apache.beam.sdk.coders.IterableCoder) IterableLikeCoder(org.apache.beam.sdk.coders.IterableLikeCoder) ModelCoderRegistrar(org.apache.beam.runners.core.construction.ModelCoderRegistrar) BeamFnStateClient(org.apache.beam.fn.harness.state.BeamFnStateClient) ArrayList(java.util.ArrayList) StateResponse(org.apache.beam.model.fnexecution.v1.BeamFnApi.StateResponse) TranslationContext(org.apache.beam.runners.core.construction.CoderTranslation.TranslationContext)

Example 35 with List

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

the class GrpcDataServiceTest method testMultipleClientsSendMessagesAreDirectedToProperConsumers.

@Test
public void testMultipleClientsSendMessagesAreDirectedToProperConsumers() throws Exception {
    final LinkedBlockingQueue<BeamFnApi.Elements> clientInboundElements = new LinkedBlockingQueue<>();
    ExecutorService executorService = Executors.newCachedThreadPool();
    final CountDownLatch waitForInboundElements = new CountDownLatch(1);
    GrpcDataService service = GrpcDataService.create(PipelineOptionsFactory.create(), Executors.newCachedThreadPool(), OutboundObserverFactory.serverDirect());
    try (GrpcFnServer<GrpcDataService> server = GrpcFnServer.allocatePortAndCreateFor(service, InProcessServerFactory.create())) {
        Collection<Future<Void>> clientFutures = new ArrayList<>();
        for (int i = 0; i < 3; ++i) {
            final String instructionId = Integer.toString(i);
            clientFutures.add(executorService.submit(() -> {
                ManagedChannel channel = InProcessChannelBuilder.forName(server.getApiServiceDescriptor().getUrl()).build();
                StreamObserver<Elements> outboundObserver = BeamFnDataGrpc.newStub(channel).data(TestStreams.withOnNext(clientInboundElements::add).build());
                outboundObserver.onNext(elementsWithData(instructionId));
                waitForInboundElements.await();
                outboundObserver.onCompleted();
                return null;
            }));
        }
        List<Collection<WindowedValue<String>>> serverInboundValues = new ArrayList<>();
        Collection<InboundDataClient> readFutures = new ArrayList<>();
        for (int i = 0; i < 3; ++i) {
            final Collection<WindowedValue<String>> serverInboundValue = new ArrayList<>();
            serverInboundValues.add(serverInboundValue);
            readFutures.add(service.receive(LogicalEndpoint.data(Integer.toString(i), TRANSFORM_ID), CODER, serverInboundValue::add));
        }
        for (InboundDataClient readFuture : readFutures) {
            readFuture.awaitCompletion();
        }
        waitForInboundElements.countDown();
        for (Future<Void> clientFuture : clientFutures) {
            clientFuture.get();
        }
        for (int i = 0; i < 3; ++i) {
            assertThat(serverInboundValues.get(i), contains(WindowedValue.valueInGlobalWindow("A" + i), WindowedValue.valueInGlobalWindow("B" + i), WindowedValue.valueInGlobalWindow("C" + i)));
        }
        assertThat(clientInboundElements, empty());
    }
}
Also used : StreamObserver(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver) ArrayList(java.util.ArrayList) ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) Elements(org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements) LinkedBlockingQueue(java.util.concurrent.LinkedBlockingQueue) CountDownLatch(java.util.concurrent.CountDownLatch) LogicalEndpoint(org.apache.beam.sdk.fn.data.LogicalEndpoint) InboundDataClient(org.apache.beam.sdk.fn.data.InboundDataClient) WindowedValue(org.apache.beam.sdk.util.WindowedValue) ExecutorService(java.util.concurrent.ExecutorService) Future(java.util.concurrent.Future) ManagedChannel(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.ManagedChannel) Collection(java.util.Collection) Test(org.junit.Test)

Aggregations

ArrayList (java.util.ArrayList)27 Test (org.junit.Test)24 ByteString (org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString)15 Elements (org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements)13 List (java.util.List)9 ExecutorService (java.util.concurrent.ExecutorService)9 WindowedValue (org.apache.beam.sdk.util.WindowedValue)9 IOException (java.io.IOException)8 Collection (java.util.Collection)8 CountDownLatch (java.util.concurrent.CountDownLatch)8 BeamFnApi (org.apache.beam.model.fnexecution.v1.BeamFnApi)8 Coder (org.apache.beam.sdk.coders.Coder)8 KvCoder (org.apache.beam.sdk.coders.KvCoder)8 Map (java.util.Map)7 ProcessBundleDescriptor (org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleDescriptor)7 RunnerApi (org.apache.beam.model.pipeline.v1.RunnerApi)7 StringUtf8Coder (org.apache.beam.sdk.coders.StringUtf8Coder)7 Iterables (org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables)7 Collections (java.util.Collections)6 CompletableFuture (java.util.concurrent.CompletableFuture)6