Search in sources :

Example 21 with KvCoder

use of org.apache.beam.sdk.coders.KvCoder in project beam by apache.

the class FlinkStreamingPortablePipelineTranslator method addGBK.

private <K, V> SingleOutputStreamOperator<WindowedValue<KV<K, Iterable<V>>>> addGBK(DataStream<WindowedValue<KV<K, V>>> inputDataStream, WindowingStrategy<?, ?> windowingStrategy, WindowedValueCoder<KV<K, V>> windowedInputCoder, String operatorName, StreamingTranslationContext context) {
    KvCoder<K, V> inputElementCoder = (KvCoder<K, V>) windowedInputCoder.getValueCoder();
    SingletonKeyedWorkItemCoder<K, V> workItemCoder = SingletonKeyedWorkItemCoder.of(inputElementCoder.getKeyCoder(), inputElementCoder.getValueCoder(), windowingStrategy.getWindowFn().windowCoder());
    WindowedValue.FullWindowedValueCoder<KeyedWorkItem<K, V>> windowedWorkItemCoder = WindowedValue.getFullCoder(workItemCoder, windowingStrategy.getWindowFn().windowCoder());
    CoderTypeInformation<WindowedValue<KeyedWorkItem<K, V>>> workItemTypeInfo = new CoderTypeInformation<>(windowedWorkItemCoder, context.getPipelineOptions());
    DataStream<WindowedValue<KeyedWorkItem<K, V>>> workItemStream = inputDataStream.flatMap(new FlinkStreamingTransformTranslators.ToKeyedWorkItem<>(context.getPipelineOptions())).returns(workItemTypeInfo).name("ToKeyedWorkItem");
    WorkItemKeySelector<K, V> keySelector = new WorkItemKeySelector<>(inputElementCoder.getKeyCoder(), new SerializablePipelineOptions(context.getPipelineOptions()));
    KeyedStream<WindowedValue<KeyedWorkItem<K, V>>, ByteBuffer> keyedWorkItemStream = workItemStream.keyBy(keySelector);
    SystemReduceFn<K, V, Iterable<V>, Iterable<V>, BoundedWindow> reduceFn = SystemReduceFn.buffering(inputElementCoder.getValueCoder());
    Coder<Iterable<V>> accumulatorCoder = IterableCoder.of(inputElementCoder.getValueCoder());
    Coder<WindowedValue<KV<K, Iterable<V>>>> outputCoder = WindowedValue.getFullCoder(KvCoder.of(inputElementCoder.getKeyCoder(), accumulatorCoder), windowingStrategy.getWindowFn().windowCoder());
    TypeInformation<WindowedValue<KV<K, Iterable<V>>>> outputTypeInfo = new CoderTypeInformation<>(outputCoder, context.getPipelineOptions());
    TupleTag<KV<K, Iterable<V>>> mainTag = new TupleTag<>("main output");
    WindowDoFnOperator<K, V, Iterable<V>> doFnOperator = new WindowDoFnOperator<>(reduceFn, operatorName, windowedWorkItemCoder, mainTag, Collections.emptyList(), new DoFnOperator.MultiOutputOutputManagerFactory<>(mainTag, outputCoder, new SerializablePipelineOptions(context.getPipelineOptions())), windowingStrategy, new HashMap<>(), /* side-input mapping */
    Collections.emptyList(), /* side inputs */
    context.getPipelineOptions(), inputElementCoder.getKeyCoder(), keySelector);
    return keyedWorkItemStream.transform(operatorName, outputTypeInfo, doFnOperator);
}
Also used : WorkItemKeySelector(org.apache.beam.runners.flink.translation.wrappers.streaming.WorkItemKeySelector) TupleTag(org.apache.beam.sdk.values.TupleTag) WindowDoFnOperator(org.apache.beam.runners.flink.translation.wrappers.streaming.WindowDoFnOperator) DoFnOperator(org.apache.beam.runners.flink.translation.wrappers.streaming.DoFnOperator) ExecutableStageDoFnOperator(org.apache.beam.runners.flink.translation.wrappers.streaming.ExecutableStageDoFnOperator) KV(org.apache.beam.sdk.values.KV) WindowedValue(org.apache.beam.sdk.util.WindowedValue) BoundedWindow(org.apache.beam.sdk.transforms.windowing.BoundedWindow) SerializablePipelineOptions(org.apache.beam.runners.core.construction.SerializablePipelineOptions) CoderTypeInformation(org.apache.beam.runners.flink.translation.types.CoderTypeInformation) WindowDoFnOperator(org.apache.beam.runners.flink.translation.wrappers.streaming.WindowDoFnOperator) KvCoder(org.apache.beam.sdk.coders.KvCoder) KV(org.apache.beam.sdk.values.KV) KeyedWorkItem(org.apache.beam.runners.core.KeyedWorkItem) ByteBuffer(java.nio.ByteBuffer)

Example 22 with KvCoder

use of org.apache.beam.sdk.coders.KvCoder in project beam by apache.

the class FlinkStreamingPortablePipelineTranslator method transformSideInputs.

private TransformedSideInputs transformSideInputs(RunnerApi.ExecutableStagePayload stagePayload, RunnerApi.Components components, StreamingTranslationContext context) {
    LinkedHashMap<RunnerApi.ExecutableStagePayload.SideInputId, PCollectionView<?>> sideInputs = getSideInputIdToPCollectionViewMap(stagePayload, components);
    Map<TupleTag<?>, Integer> tagToIntMapping = new HashMap<>();
    Map<Integer, PCollectionView<?>> intToViewMapping = new HashMap<>();
    List<WindowedValueCoder<KV<Void, Object>>> kvCoders = new ArrayList<>();
    List<Coder<?>> viewCoders = new ArrayList<>();
    int count = 0;
    for (Map.Entry<RunnerApi.ExecutableStagePayload.SideInputId, PCollectionView<?>> sideInput : sideInputs.entrySet()) {
        TupleTag<?> tag = sideInput.getValue().getTagInternal();
        intToViewMapping.put(count, sideInput.getValue());
        tagToIntMapping.put(tag, count);
        count++;
        String collectionId = components.getTransformsOrThrow(sideInput.getKey().getTransformId()).getInputsOrThrow(sideInput.getKey().getLocalName());
        DataStream<Object> sideInputStream = context.getDataStreamOrThrow(collectionId);
        TypeInformation<Object> tpe = sideInputStream.getType();
        if (!(tpe instanceof CoderTypeInformation)) {
            throw new IllegalStateException("Input Stream TypeInformation is no CoderTypeInformation.");
        }
        WindowedValueCoder<Object> coder = (WindowedValueCoder) ((CoderTypeInformation) tpe).getCoder();
        Coder<KV<Void, Object>> kvCoder = KvCoder.of(VoidCoder.of(), coder.getValueCoder());
        kvCoders.add(coder.withValueCoder(kvCoder));
        // coder for materialized view matching GBK below
        WindowedValueCoder<KV<Void, Iterable<Object>>> viewCoder = coder.withValueCoder(KvCoder.of(VoidCoder.of(), IterableCoder.of(coder.getValueCoder())));
        viewCoders.add(viewCoder);
    }
    // second pass, now that we gathered the input coders
    UnionCoder unionCoder = UnionCoder.of(viewCoders);
    CoderTypeInformation<RawUnionValue> unionTypeInformation = new CoderTypeInformation<>(unionCoder, context.getPipelineOptions());
    // transform each side input to RawUnionValue and union them
    DataStream<RawUnionValue> sideInputUnion = null;
    for (Map.Entry<RunnerApi.ExecutableStagePayload.SideInputId, PCollectionView<?>> sideInput : sideInputs.entrySet()) {
        TupleTag<?> tag = sideInput.getValue().getTagInternal();
        final int intTag = tagToIntMapping.get(tag);
        RunnerApi.PTransform pTransform = components.getTransformsOrThrow(sideInput.getKey().getTransformId());
        String collectionId = pTransform.getInputsOrThrow(sideInput.getKey().getLocalName());
        DataStream<WindowedValue<?>> sideInputStream = context.getDataStreamOrThrow(collectionId);
        // insert GBK to materialize side input view
        String viewName = sideInput.getKey().getTransformId() + "-" + sideInput.getKey().getLocalName();
        WindowedValueCoder<KV<Void, Object>> kvCoder = kvCoders.get(intTag);
        DataStream<WindowedValue<KV<Void, Object>>> keyedSideInputStream = sideInputStream.map(new ToVoidKeyValue(context.getPipelineOptions()));
        SingleOutputStreamOperator<WindowedValue<KV<Void, Iterable<Object>>>> viewStream = addGBK(keyedSideInputStream, sideInput.getValue().getWindowingStrategyInternal(), kvCoder, viewName, context);
        // Assign a unique but consistent id to re-map operator state
        viewStream.uid(pTransform.getUniqueName() + "-" + sideInput.getKey().getLocalName());
        DataStream<RawUnionValue> unionValueStream = viewStream.map(new FlinkStreamingTransformTranslators.ToRawUnion<>(intTag, context.getPipelineOptions())).returns(unionTypeInformation);
        if (sideInputUnion == null) {
            sideInputUnion = unionValueStream;
        } else {
            sideInputUnion = sideInputUnion.union(unionValueStream);
        }
    }
    return new TransformedSideInputs(intToViewMapping, sideInputUnion);
}
Also used : LinkedHashMap(java.util.LinkedHashMap) HashMap(java.util.HashMap) ArrayList(java.util.ArrayList) TupleTag(org.apache.beam.sdk.values.TupleTag) RunnerApi(org.apache.beam.model.pipeline.v1.RunnerApi) WindowedValue(org.apache.beam.sdk.util.WindowedValue) CoderTypeInformation(org.apache.beam.runners.flink.translation.types.CoderTypeInformation) SingletonKeyedWorkItemCoder(org.apache.beam.runners.flink.translation.wrappers.streaming.SingletonKeyedWorkItemCoder) WindowedValueCoder(org.apache.beam.sdk.util.WindowedValue.WindowedValueCoder) KvCoder(org.apache.beam.sdk.coders.KvCoder) PipelineTranslatorUtils.instantiateCoder(org.apache.beam.runners.fnexecution.translation.PipelineTranslatorUtils.instantiateCoder) IterableCoder(org.apache.beam.sdk.coders.IterableCoder) VoidCoder(org.apache.beam.sdk.coders.VoidCoder) UnionCoder(org.apache.beam.sdk.transforms.join.UnionCoder) Coder(org.apache.beam.sdk.coders.Coder) ByteArrayCoder(org.apache.beam.sdk.coders.ByteArrayCoder) UnionCoder(org.apache.beam.sdk.transforms.join.UnionCoder) RawUnionValue(org.apache.beam.sdk.transforms.join.RawUnionValue) KV(org.apache.beam.sdk.values.KV) RunnerPCollectionView(org.apache.beam.runners.core.construction.RunnerPCollectionView) PCollectionView(org.apache.beam.sdk.values.PCollectionView) WindowedValueCoder(org.apache.beam.sdk.util.WindowedValue.WindowedValueCoder) ImmutableMap(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap) Map(java.util.Map) LinkedHashMap(java.util.LinkedHashMap) BiMap(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.BiMap) TreeMap(java.util.TreeMap) PipelineTranslatorUtils.createOutputMap(org.apache.beam.runners.fnexecution.translation.PipelineTranslatorUtils.createOutputMap) HashMap(java.util.HashMap)

Example 23 with KvCoder

use of org.apache.beam.sdk.coders.KvCoder in project beam by apache.

the class SideInputInitializer method initializeBroadcastVariable.

@Override
public Map<BoundedWindow, ViewT> initializeBroadcastVariable(Iterable<WindowedValue<?>> inputValues) {
    // first partition into windows
    Map<BoundedWindow, List<WindowedValue<?>>> partitionedElements = new HashMap<>();
    for (WindowedValue<?> value : inputValues) {
        for (BoundedWindow window : value.getWindows()) {
            List<WindowedValue<?>> windowedValues = partitionedElements.computeIfAbsent(window, k -> new ArrayList<>());
            windowedValues.add(value);
        }
    }
    Map<BoundedWindow, ViewT> resultMap = new HashMap<>();
    for (Map.Entry<BoundedWindow, List<WindowedValue<?>>> elements : partitionedElements.entrySet()) {
        switch(view.getViewFn().getMaterialization().getUrn()) {
            case Materializations.ITERABLE_MATERIALIZATION_URN:
                {
                    ViewFn<IterableView, ViewT> viewFn = (ViewFn<IterableView, ViewT>) view.getViewFn();
                    resultMap.put(elements.getKey(), viewFn.apply(() -> elements.getValue().stream().map(WindowedValue::getValue).collect(Collectors.toList())));
                }
                break;
            case Materializations.MULTIMAP_MATERIALIZATION_URN:
                {
                    ViewFn<MultimapView, ViewT> viewFn = (ViewFn<MultimapView, ViewT>) view.getViewFn();
                    Coder<?> keyCoder = ((KvCoder<?, ?>) view.getCoderInternal()).getKeyCoder();
                    resultMap.put(elements.getKey(), viewFn.apply(InMemoryMultimapSideInputView.fromIterable(keyCoder, (Iterable) elements.getValue().stream().map(WindowedValue::getValue).collect(Collectors.toList()))));
                }
                break;
            default:
                throw new IllegalStateException(String.format("Unknown side input materialization format requested '%s'", view.getViewFn().getMaterialization().getUrn()));
        }
    }
    return resultMap;
}
Also used : KvCoder(org.apache.beam.sdk.coders.KvCoder) Coder(org.apache.beam.sdk.coders.Coder) IterableView(org.apache.beam.sdk.transforms.Materializations.IterableView) HashMap(java.util.HashMap) MultimapView(org.apache.beam.sdk.transforms.Materializations.MultimapView) ViewFn(org.apache.beam.sdk.transforms.ViewFn) WindowedValue(org.apache.beam.sdk.util.WindowedValue) BoundedWindow(org.apache.beam.sdk.transforms.windowing.BoundedWindow) ArrayList(java.util.ArrayList) List(java.util.List) HashMap(java.util.HashMap) Map(java.util.Map)

Example 24 with KvCoder

use of org.apache.beam.sdk.coders.KvCoder in project beam by apache.

the class BeamFnMapTaskExecutorFactory method createPartialGroupByKeyOperation.

<K> OperationNode createPartialGroupByKeyOperation(Network<Node, Edge> network, ParallelInstructionNode node, PipelineOptions options, DataflowExecutionContext<?> executionContext, DataflowOperationContext operationContext) throws Exception {
    ParallelInstruction instruction = node.getParallelInstruction();
    PartialGroupByKeyInstruction pgbk = instruction.getPartialGroupByKey();
    OutputReceiver[] receivers = getOutputReceivers(network, node);
    Coder<?> windowedCoder = CloudObjects.coderFromCloudObject(CloudObject.fromSpec(pgbk.getInputElementCodec()));
    if (!(windowedCoder instanceof WindowedValueCoder)) {
        throw new IllegalArgumentException(String.format("unexpected kind of input coder for PartialGroupByKeyOperation: %s", windowedCoder));
    }
    Coder<?> elemCoder = ((WindowedValueCoder<?>) windowedCoder).getValueCoder();
    if (!(elemCoder instanceof KvCoder)) {
        throw new IllegalArgumentException(String.format("unexpected kind of input element coder for PartialGroupByKeyOperation: %s", elemCoder));
    }
    @SuppressWarnings("unchecked") KvCoder<K, ?> keyedElementCoder = (KvCoder<K, ?>) elemCoder;
    CloudObject cloudUserFn = pgbk.getValueCombiningFn() != null ? CloudObject.fromSpec(pgbk.getValueCombiningFn()) : null;
    ParDoFn fn = PartialGroupByKeyParDoFns.create(options, keyedElementCoder, cloudUserFn, pgbk.getSideInputs(), Arrays.<Receiver>asList(receivers), executionContext, operationContext);
    return OperationNode.create(new ParDoOperation(fn, receivers, operationContext));
}
Also used : OutputReceiver(org.apache.beam.runners.dataflow.worker.util.common.worker.OutputReceiver) KvCoder(org.apache.beam.sdk.coders.KvCoder) ParDoFn(org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoFn) ParDoOperation(org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoOperation) ParallelInstruction(com.google.api.services.dataflow.model.ParallelInstruction) WindowedValueCoder(org.apache.beam.sdk.util.WindowedValue.WindowedValueCoder) CloudObject(org.apache.beam.runners.dataflow.util.CloudObject) PartialGroupByKeyInstruction(com.google.api.services.dataflow.model.PartialGroupByKeyInstruction)

Example 25 with KvCoder

use of org.apache.beam.sdk.coders.KvCoder in project beam by apache.

the class StreamingDataflowWorker method process.

private void process(final SdkWorkerHarness worker, final ComputationState computationState, final Instant inputDataWatermark, @Nullable final Instant outputDataWatermark, @Nullable final Instant synchronizedProcessingTime, final Work work) {
    final Windmill.WorkItem workItem = work.getWorkItem();
    final String computationId = computationState.getComputationId();
    final ByteString key = workItem.getKey();
    work.setState(State.PROCESSING);
    {
        StringBuilder workIdBuilder = new StringBuilder(33);
        workIdBuilder.append(Long.toHexString(workItem.getShardingKey()));
        workIdBuilder.append('-');
        workIdBuilder.append(Long.toHexString(workItem.getWorkToken()));
        DataflowWorkerLoggingMDC.setWorkId(workIdBuilder.toString());
    }
    DataflowWorkerLoggingMDC.setStageName(computationId);
    LOG.debug("Starting processing for {}:\n{}", computationId, work);
    Windmill.WorkItemCommitRequest.Builder outputBuilder = initializeOutputBuilder(key, workItem);
    // Before any processing starts, call any pending OnCommit callbacks.  Nothing that requires
    // cleanup should be done before this, since we might exit early here.
    callFinalizeCallbacks(workItem);
    if (workItem.getSourceState().getOnlyFinalize()) {
        outputBuilder.setSourceStateUpdates(Windmill.SourceState.newBuilder().setOnlyFinalize(true));
        work.setState(State.COMMIT_QUEUED);
        commitQueue.put(new Commit(outputBuilder.build(), computationState, work));
        return;
    }
    long processingStartTimeNanos = System.nanoTime();
    final MapTask mapTask = computationState.getMapTask();
    StageInfo stageInfo = stageInfoMap.computeIfAbsent(mapTask.getStageName(), s -> new StageInfo(s, mapTask.getSystemName(), this));
    ExecutionState executionState = null;
    try {
        executionState = computationState.getExecutionStateQueue(worker).poll();
        if (executionState == null) {
            MutableNetwork<Node, Edge> mapTaskNetwork = mapTaskToNetwork.apply(mapTask);
            if (LOG.isDebugEnabled()) {
                LOG.debug("Network as Graphviz .dot: {}", Networks.toDot(mapTaskNetwork));
            }
            ParallelInstructionNode readNode = (ParallelInstructionNode) Iterables.find(mapTaskNetwork.nodes(), node -> node instanceof ParallelInstructionNode && ((ParallelInstructionNode) node).getParallelInstruction().getRead() != null);
            InstructionOutputNode readOutputNode = (InstructionOutputNode) Iterables.getOnlyElement(mapTaskNetwork.successors(readNode));
            DataflowExecutionContext.DataflowExecutionStateTracker executionStateTracker = new DataflowExecutionContext.DataflowExecutionStateTracker(ExecutionStateSampler.instance(), stageInfo.executionStateRegistry.getState(NameContext.forStage(mapTask.getStageName()), "other", null, ScopedProfiler.INSTANCE.emptyScope()), stageInfo.deltaCounters, options, computationId);
            StreamingModeExecutionContext context = new StreamingModeExecutionContext(pendingDeltaCounters, computationId, readerCache, !computationState.getTransformUserNameToStateFamily().isEmpty() ? computationState.getTransformUserNameToStateFamily() : stateNameMap, stateCache.forComputation(computationId), stageInfo.metricsContainerRegistry, executionStateTracker, stageInfo.executionStateRegistry, maxSinkBytes);
            DataflowMapTaskExecutor mapTaskExecutor = mapTaskExecutorFactory.create(worker.getControlClientHandler(), worker.getGrpcDataFnServer(), sdkHarnessRegistry.beamFnDataApiServiceDescriptor(), worker.getGrpcStateFnServer(), mapTaskNetwork, options, mapTask.getStageName(), readerRegistry, sinkRegistry, context, pendingDeltaCounters, idGenerator);
            ReadOperation readOperation = mapTaskExecutor.getReadOperation();
            // Disable progress updates since its results are unused  for streaming
            // and involves starting a thread.
            readOperation.setProgressUpdatePeriodMs(ReadOperation.DONT_UPDATE_PERIODICALLY);
            Preconditions.checkState(mapTaskExecutor.supportsRestart(), "Streaming runner requires all operations support restart.");
            Coder<?> readCoder;
            readCoder = CloudObjects.coderFromCloudObject(CloudObject.fromSpec(readOutputNode.getInstructionOutput().getCodec()));
            Coder<?> keyCoder = extractKeyCoder(readCoder);
            // If using a custom source, count bytes read for autoscaling.
            if (CustomSources.class.getName().equals(readNode.getParallelInstruction().getRead().getSource().getSpec().get("@type"))) {
                NameContext nameContext = NameContext.create(mapTask.getStageName(), readNode.getParallelInstruction().getOriginalName(), readNode.getParallelInstruction().getSystemName(), readNode.getParallelInstruction().getName());
                readOperation.receivers[0].addOutputCounter(new OutputObjectAndByteCounter(new IntrinsicMapTaskExecutorFactory.ElementByteSizeObservableCoder<>(readCoder), mapTaskExecutor.getOutputCounters(), nameContext).setSamplingPeriod(100).countBytes("dataflow_input_size-" + mapTask.getSystemName()));
            }
            executionState = new ExecutionState(mapTaskExecutor, context, keyCoder, executionStateTracker);
        }
        WindmillStateReader stateReader = new WindmillStateReader(metricTrackingWindmillServer, computationId, key, workItem.getShardingKey(), workItem.getWorkToken());
        StateFetcher localStateFetcher = stateFetcher.byteTrackingView();
        // If the read output KVs, then we can decode Windmill's byte key into a userland
        // key object and provide it to the execution context for use with per-key state.
        // Otherwise, we pass null.
        // 
        // The coder type that will be present is:
        // WindowedValueCoder(TimerOrElementCoder(KvCoder))
        @Nullable Coder<?> keyCoder = executionState.getKeyCoder();
        @Nullable Object executionKey = keyCoder == null ? null : keyCoder.decode(key.newInput(), Coder.Context.OUTER);
        if (workItem.hasHotKeyInfo()) {
            Windmill.HotKeyInfo hotKeyInfo = workItem.getHotKeyInfo();
            Duration hotKeyAge = Duration.millis(hotKeyInfo.getHotKeyAgeUsec() / 1000);
            // The MapTask instruction is ordered by dependencies, such that the first element is
            // always going to be the shuffle task.
            String stepName = computationState.getMapTask().getInstructions().get(0).getName();
            if (options.isHotKeyLoggingEnabled() && keyCoder != null) {
                hotKeyLogger.logHotKeyDetection(stepName, hotKeyAge, executionKey);
            } else {
                hotKeyLogger.logHotKeyDetection(stepName, hotKeyAge);
            }
        }
        executionState.getContext().start(executionKey, workItem, inputDataWatermark, outputDataWatermark, synchronizedProcessingTime, stateReader, localStateFetcher, outputBuilder);
        // Blocks while executing work.
        executionState.getWorkExecutor().execute();
        Iterables.addAll(this.pendingMonitoringInfos, executionState.getWorkExecutor().extractMetricUpdates());
        commitCallbacks.putAll(executionState.getContext().flushState());
        // Release the execution state for another thread to use.
        computationState.getExecutionStateQueue(worker).offer(executionState);
        executionState = null;
        // Add the output to the commit queue.
        work.setState(State.COMMIT_QUEUED);
        WorkItemCommitRequest commitRequest = outputBuilder.build();
        int byteLimit = maxWorkItemCommitBytes;
        int commitSize = commitRequest.getSerializedSize();
        int estimatedCommitSize = commitSize < 0 ? Integer.MAX_VALUE : commitSize;
        // Detect overflow of integer serialized size or if the byte limit was exceeded.
        windmillMaxObservedWorkItemCommitBytes.addValue(estimatedCommitSize);
        if (commitSize < 0 || commitSize > byteLimit) {
            KeyCommitTooLargeException e = KeyCommitTooLargeException.causedBy(computationId, byteLimit, commitRequest);
            reportFailure(computationId, workItem, e);
            LOG.error(e.toString());
            // Drop the current request in favor of a new, minimal one requesting truncation.
            // Messages, timers, counters, and other commit content will not be used by the service
            // so we're purposefully dropping them here
            commitRequest = buildWorkItemTruncationRequest(key, workItem, estimatedCommitSize);
        }
        commitQueue.put(new Commit(commitRequest, computationState, work));
        // Compute shuffle and state byte statistics these will be flushed asynchronously.
        long stateBytesWritten = outputBuilder.clearOutputMessages().build().getSerializedSize();
        long shuffleBytesRead = 0;
        for (Windmill.InputMessageBundle bundle : workItem.getMessageBundlesList()) {
            for (Windmill.Message message : bundle.getMessagesList()) {
                shuffleBytesRead += message.getSerializedSize();
            }
        }
        long stateBytesRead = stateReader.getBytesRead() + localStateFetcher.getBytesRead();
        windmillShuffleBytesRead.addValue(shuffleBytesRead);
        windmillStateBytesRead.addValue(stateBytesRead);
        windmillStateBytesWritten.addValue(stateBytesWritten);
        LOG.debug("Processing done for work token: {}", workItem.getWorkToken());
    } catch (Throwable t) {
        if (executionState != null) {
            try {
                executionState.getContext().invalidateCache();
                executionState.getWorkExecutor().close();
            } catch (Exception e) {
                LOG.warn("Failed to close map task executor: ", e);
            } finally {
                // Release references to potentially large objects early.
                executionState = null;
            }
        }
        t = t instanceof UserCodeException ? t.getCause() : t;
        boolean retryLocally = false;
        if (KeyTokenInvalidException.isKeyTokenInvalidException(t)) {
            LOG.debug("Execution of work for computation '{}' on key '{}' failed due to token expiration. " + "Work will not be retried locally.", computationId, key.toStringUtf8());
        } else {
            LastExceptionDataProvider.reportException(t);
            LOG.debug("Failed work: {}", work);
            Duration elapsedTimeSinceStart = new Duration(Instant.now(), work.getStartTime());
            if (!reportFailure(computationId, workItem, t)) {
                LOG.error("Execution of work for computation '{}' on key '{}' failed with uncaught exception, " + "and Windmill indicated not to retry locally.", computationId, key.toStringUtf8(), t);
            } else if (isOutOfMemoryError(t)) {
                File heapDump = memoryMonitor.tryToDumpHeap();
                LOG.error("Execution of work for computation '{}' for key '{}' failed with out-of-memory. " + "Work will not be retried locally. Heap dump {}.", computationId, key.toStringUtf8(), heapDump == null ? "not written" : ("written to '" + heapDump + "'"), t);
            } else if (elapsedTimeSinceStart.isLongerThan(MAX_LOCAL_PROCESSING_RETRY_DURATION)) {
                LOG.error("Execution of work for computation '{}' for key '{}' failed with uncaught exception, " + "and it will not be retried locally because the elapsed time since start {} " + "exceeds {}.", computationId, key.toStringUtf8(), elapsedTimeSinceStart, MAX_LOCAL_PROCESSING_RETRY_DURATION, t);
            } else {
                LOG.error("Execution of work for computation '{}' on key '{}' failed with uncaught exception. " + "Work will be retried locally.", computationId, key.toStringUtf8(), t);
                retryLocally = true;
            }
        }
        if (retryLocally) {
            // Try again after some delay and at the end of the queue to avoid a tight loop.
            sleep(retryLocallyDelayMs);
            workUnitExecutor.forceExecute(work, work.getWorkItem().getSerializedSize());
        } else {
            // Consider the item invalid. It will eventually be retried by Windmill if it still needs to
            // be processed.
            computationState.completeWork(ShardedKey.create(key, workItem.getShardingKey()), workItem.getWorkToken());
        }
    } finally {
        // Update total processing time counters. Updating in finally clause ensures that
        // work items causing exceptions are also accounted in time spent.
        long processingTimeMsecs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - processingStartTimeNanos);
        stageInfo.totalProcessingMsecs.addValue(processingTimeMsecs);
        // either here or in DFE.
        if (work.getWorkItem().hasTimers()) {
            stageInfo.timerProcessingMsecs.addValue(processingTimeMsecs);
        }
        DataflowWorkerLoggingMDC.setWorkId(null);
        DataflowWorkerLoggingMDC.setStageName(null);
    }
}
Also used : MetricName(org.apache.beam.sdk.metrics.MetricName) MapTask(com.google.api.services.dataflow.model.MapTask) UserCodeException(org.apache.beam.sdk.util.UserCodeException) WindowedValueCoder(org.apache.beam.sdk.util.WindowedValue.WindowedValueCoder) MetricsLogger(org.apache.beam.runners.core.metrics.MetricsLogger) CommitWorkStream(org.apache.beam.runners.dataflow.worker.windmill.WindmillServerStub.CommitWorkStream) CloudObjects(org.apache.beam.runners.dataflow.util.CloudObjects) ImmutableMap(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap) Map(java.util.Map) CreateRegisterFnOperationFunction(org.apache.beam.runners.dataflow.worker.graph.CreateRegisterFnOperationFunction) ScopedProfiler(org.apache.beam.runners.dataflow.worker.profiler.ScopedProfiler) StreamPool(org.apache.beam.runners.dataflow.worker.windmill.WindmillServerStub.StreamPool) DataflowCounterUpdateExtractor(org.apache.beam.runners.dataflow.worker.counters.DataflowCounterUpdateExtractor) Uninterruptibles(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.Uninterruptibles) TimerTask(java.util.TimerTask) WorkItemStatus(com.google.api.services.dataflow.model.WorkItemStatus) WorkerStatusPages(org.apache.beam.runners.dataflow.worker.status.WorkerStatusPages) RegisterNodeFunction(org.apache.beam.runners.dataflow.worker.graph.RegisterNodeFunction) IdGenerator(org.apache.beam.sdk.fn.IdGenerator) PrintWriter(java.io.PrintWriter) KvCoder(org.apache.beam.sdk.coders.KvCoder) THROTTLING_MSECS_METRIC_NAME(org.apache.beam.runners.dataflow.worker.DataflowSystemMetrics.THROTTLING_MSECS_METRIC_NAME) ReadOperation(org.apache.beam.runners.dataflow.worker.util.common.worker.ReadOperation) CacheBuilder(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheBuilder) Sleeper(org.apache.beam.sdk.util.Sleeper) StreamingModeExecutionStateRegistry(org.apache.beam.runners.dataflow.worker.StreamingModeExecutionContext.StreamingModeExecutionStateRegistry) DebugCapture(org.apache.beam.runners.dataflow.worker.status.DebugCapture) Executors(java.util.concurrent.Executors) BoundedQueueExecutor(org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor) MultimapBuilder(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.MultimapBuilder) VisibleForTesting(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting) WorkItemCommitRequest(org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItemCommitRequest) AutoValue(com.google.auto.value.AutoValue) Counter(org.apache.beam.runners.dataflow.worker.counters.Counter) InsertFetchAndFilterStreamingSideInputNodes(org.apache.beam.runners.dataflow.worker.graph.InsertFetchAndFilterStreamingSideInputNodes) Capturable(org.apache.beam.runners.dataflow.worker.status.DebugCapture.Capturable) Networks(org.apache.beam.runners.dataflow.worker.graph.Networks) DeduceNodeLocationsFunction(org.apache.beam.runners.dataflow.worker.graph.DeduceNodeLocationsFunction) ExecutionStateTracker(org.apache.beam.runners.core.metrics.ExecutionStateTracker) Cache(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.Cache) Duration(org.joda.time.Duration) Splitter(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Splitter) Optional(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Optional) ArrayList(java.util.ArrayList) CounterSet(org.apache.beam.runners.dataflow.worker.counters.CounterSet) Status(com.google.api.services.dataflow.model.Status) HttpServletRequest(javax.servlet.http.HttpServletRequest) EvictingQueue(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.EvictingQueue) GetWorkStream(org.apache.beam.runners.dataflow.worker.windmill.WindmillServerStub.GetWorkStream) Preconditions.checkArgument(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument) RunnerApi(org.apache.beam.model.pipeline.v1.RunnerApi) Windmill(org.apache.beam.runners.dataflow.worker.windmill.Windmill) StreamingComputationConfig(com.google.api.services.dataflow.model.StreamingComputationConfig) DataflowRunner(org.apache.beam.runners.dataflow.DataflowRunner) IOException(java.io.IOException) File(java.io.File) ExecutionException(java.util.concurrent.ExecutionException) CounterStructuredName(com.google.api.services.dataflow.model.CounterStructuredName) AtomicLong(java.util.concurrent.atomic.AtomicLong) MetricsEnvironment(org.apache.beam.sdk.metrics.MetricsEnvironment) DataflowWorkerLoggingMDC(org.apache.beam.runners.dataflow.worker.logging.DataflowWorkerLoggingMDC) RemoteGrpcPortNode(org.apache.beam.runners.dataflow.worker.graph.Nodes.RemoteGrpcPortNode) ArrayDeque(java.util.ArrayDeque) FileSystems(org.apache.beam.sdk.io.FileSystems) StreamingPerStageSystemCounterNames(org.apache.beam.runners.dataflow.worker.DataflowSystemMetrics.StreamingPerStageSystemCounterNames) State(org.apache.beam.runners.dataflow.worker.StreamingDataflowWorker.Work.State) CounterUpdateAggregators(org.apache.beam.runners.dataflow.worker.counters.CounterUpdateAggregators) Edge(org.apache.beam.runners.dataflow.worker.graph.Edges.Edge) ReplacePgbkWithPrecombineFunction(org.apache.beam.runners.dataflow.worker.graph.ReplacePgbkWithPrecombineFunction) OutputObjectAndByteCounter(org.apache.beam.runners.dataflow.worker.util.common.worker.OutputObjectAndByteCounter) MoreObjects(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.MoreObjects) LoggerFactory(org.slf4j.LoggerFactory) Random(java.util.Random) Timer(java.util.Timer) MutableNetwork(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.graph.MutableNetwork) HostAndPort(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.net.HostAndPort) BackOff(org.apache.beam.sdk.util.BackOff) BackOffUtils(org.apache.beam.sdk.util.BackOffUtils) StatusDataProvider(org.apache.beam.runners.dataflow.worker.status.StatusDataProvider) DataflowWorkerHarnessOptions(org.apache.beam.runners.dataflow.options.DataflowWorkerHarnessOptions) DataflowRunner.hasExperiment(org.apache.beam.runners.dataflow.DataflowRunner.hasExperiment) Transport(org.apache.beam.sdk.extensions.gcp.util.Transport) Iterables(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables) NameContext(org.apache.beam.runners.dataflow.worker.counters.NameContext) ThreadFactory(java.util.concurrent.ThreadFactory) JvmInitializers(org.apache.beam.sdk.fn.JvmInitializers) Collection(java.util.Collection) ConcurrentHashMap(java.util.concurrent.ConcurrentHashMap) SdkWorkerHarness(org.apache.beam.runners.dataflow.worker.SdkHarnessRegistry.SdkWorkerHarness) FixMultiOutputInfosOnParDoInstructions(org.apache.beam.runners.dataflow.worker.apiary.FixMultiOutputInfosOnParDoInstructions) LinkedBlockingQueue(java.util.concurrent.LinkedBlockingQueue) LastExceptionDataProvider(org.apache.beam.runners.dataflow.worker.status.LastExceptionDataProvider) List(java.util.List) ListMultimap(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ListMultimap) Queue(java.util.Queue) BaseStatusServlet(org.apache.beam.runners.dataflow.worker.status.BaseStatusServlet) SuppressFBWarnings(edu.umd.cs.findbugs.annotations.SuppressFBWarnings) ConcurrentLinkedQueue(java.util.concurrent.ConcurrentLinkedQueue) StreamingSystemCounterNames(org.apache.beam.runners.dataflow.worker.DataflowSystemMetrics.StreamingSystemCounterNames) IdGenerators(org.apache.beam.sdk.fn.IdGenerators) CustomSources(org.apache.beam.runners.dataflow.internal.CustomSources) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) Coder(org.apache.beam.sdk.coders.Coder) HashMap(java.util.HashMap) CompletableFuture(java.util.concurrent.CompletableFuture) ExecutionStateSampler(org.apache.beam.runners.core.metrics.ExecutionStateSampler) Deque(java.util.Deque) InstructionOutputNode(org.apache.beam.runners.dataflow.worker.graph.Nodes.InstructionOutputNode) WorkItem(com.google.api.services.dataflow.model.WorkItem) Function(java.util.function.Function) StreamingDataflowWorkerOptions(org.apache.beam.runners.dataflow.worker.options.StreamingDataflowWorkerOptions) TextFormat(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.TextFormat) ConcurrentMap(java.util.concurrent.ConcurrentMap) MemoryMonitor(org.apache.beam.runners.dataflow.worker.util.MemoryMonitor) HashSet(java.util.HashSet) DeduceFlattenLocationsFunction(org.apache.beam.runners.dataflow.worker.graph.DeduceFlattenLocationsFunction) StreamingConfigTask(com.google.api.services.dataflow.model.StreamingConfigTask) WindmillServerStub(org.apache.beam.runners.dataflow.worker.windmill.WindmillServerStub) ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) CloudObject(org.apache.beam.runners.dataflow.util.CloudObject) ParallelInstructionNode(org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode) Nullable(org.checkerframework.checker.nullness.qual.Nullable) CounterUpdate(com.google.api.services.dataflow.model.CounterUpdate) MapTaskToNetworkFunction(org.apache.beam.runners.dataflow.worker.graph.MapTaskToNetworkFunction) FluentBackoff(org.apache.beam.sdk.util.FluentBackoff) Logger(org.slf4j.Logger) Iterator(java.util.Iterator) CloneAmbiguousFlattensFunction(org.apache.beam.runners.dataflow.worker.graph.CloneAmbiguousFlattensFunction) UTF_8(java.nio.charset.StandardCharsets.UTF_8) Semaphore(java.util.concurrent.Semaphore) Node(org.apache.beam.runners.dataflow.worker.graph.Nodes.Node) HttpServletResponse(javax.servlet.http.HttpServletResponse) TimeUnit(java.util.concurrent.TimeUnit) Preconditions(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions) Instant(org.joda.time.Instant) RemoteGrpcPort(org.apache.beam.model.fnexecution.v1.BeamFnApi.RemoteGrpcPort) Collections(java.util.Collections) LengthPrefixUnknownCoders(org.apache.beam.runners.dataflow.worker.graph.LengthPrefixUnknownCoders) ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) RemoteGrpcPortNode(org.apache.beam.runners.dataflow.worker.graph.Nodes.RemoteGrpcPortNode) InstructionOutputNode(org.apache.beam.runners.dataflow.worker.graph.Nodes.InstructionOutputNode) ParallelInstructionNode(org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode) Node(org.apache.beam.runners.dataflow.worker.graph.Nodes.Node) ParallelInstructionNode(org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode) ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) InstructionOutputNode(org.apache.beam.runners.dataflow.worker.graph.Nodes.InstructionOutputNode) UserCodeException(org.apache.beam.sdk.util.UserCodeException) Windmill(org.apache.beam.runners.dataflow.worker.windmill.Windmill) ReadOperation(org.apache.beam.runners.dataflow.worker.util.common.worker.ReadOperation) OutputObjectAndByteCounter(org.apache.beam.runners.dataflow.worker.util.common.worker.OutputObjectAndByteCounter) CustomSources(org.apache.beam.runners.dataflow.internal.CustomSources) NameContext(org.apache.beam.runners.dataflow.worker.counters.NameContext) Duration(org.joda.time.Duration) UserCodeException(org.apache.beam.sdk.util.UserCodeException) IOException(java.io.IOException) ExecutionException(java.util.concurrent.ExecutionException) WorkItemCommitRequest(org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItemCommitRequest) MapTask(com.google.api.services.dataflow.model.MapTask) CloudObject(org.apache.beam.runners.dataflow.util.CloudObject) Edge(org.apache.beam.runners.dataflow.worker.graph.Edges.Edge) File(java.io.File) Nullable(org.checkerframework.checker.nullness.qual.Nullable)

Aggregations

KvCoder (org.apache.beam.sdk.coders.KvCoder)44 Coder (org.apache.beam.sdk.coders.Coder)26 WindowedValue (org.apache.beam.sdk.util.WindowedValue)25 KV (org.apache.beam.sdk.values.KV)21 BoundedWindow (org.apache.beam.sdk.transforms.windowing.BoundedWindow)20 Map (java.util.Map)17 List (java.util.List)16 ArrayList (java.util.ArrayList)15 RunnerApi (org.apache.beam.model.pipeline.v1.RunnerApi)15 IOException (java.io.IOException)14 HashMap (java.util.HashMap)14 WindowingStrategy (org.apache.beam.sdk.values.WindowingStrategy)13 StringUtf8Coder (org.apache.beam.sdk.coders.StringUtf8Coder)11 IterableCoder (org.apache.beam.sdk.coders.IterableCoder)10 VoidCoder (org.apache.beam.sdk.coders.VoidCoder)10 PCollectionView (org.apache.beam.sdk.values.PCollectionView)10 Test (org.junit.Test)10 WindowedValueCoder (org.apache.beam.sdk.util.WindowedValue.WindowedValueCoder)8 ViewFn (org.apache.beam.sdk.transforms.ViewFn)7 PCollection (org.apache.beam.sdk.values.PCollection)7