Examples with UserCodeException org.apache.beam.sdk.util.UserCodeException used on opensource projects

Example 1 with UserCodeException

use of org.apache.beam.sdk.util.UserCodeException in project beam by apache.

the class PipelineTest method testPipelineSDKExceptionHandling.

@Test
public void testPipelineSDKExceptionHandling() {
    PipelineOptions options = TestPipeline.testingPipelineOptions();
    options.setRunner(TestPipelineRunnerThrowingSdkException.class);
    Pipeline p = Pipeline.create(options);
    // Check pipeline runner correctly catches SDK errors.
    try {
        p.run();
        fail("Should have thrown an exception.");
    } catch (RuntimeException exn) {
        // Make sure the exception isn't a UserCodeException.
        assertThat(exn, not(instanceOf(UserCodeException.class)));
        // Assert that the message is correct.
        assertThat(exn.getMessage(), containsString("SDK exception"));
        // RuntimeException should be IllegalStateException.
        assertThat(exn, instanceOf(IllegalStateException.class));
    }
}

Example 2 with UserCodeException

use of org.apache.beam.sdk.util.UserCodeException in project beam by apache.

the class WorkItemStatusClient method reportError.

/**
 * Return the {@link WorkItemServiceState} resulting from sending an error completion status.
 */
public synchronized WorkItemServiceState reportError(Throwable e) throws IOException {
    checkState(!finalStateSent, "cannot reportUpdates after sending a final state");
    WorkItemStatus status = createStatusUpdate(true);
    // TODO: Provide more structure representation of error, e.g., the serialized exception object.
    // TODO: Look into moving the stack trace thinning into the client.
    Throwable t = e instanceof UserCodeException ? e.getCause() : e;
    Status error = new Status();
    // Code.UNKNOWN.  TODO: Replace with a generated definition.
    error.setCode(2);
    // TODO: Attach the stack trace as exception details, not to the message.
    String logPrefix = String.format("Failure processing work item %s", uniqueWorkId());
    if (isOutOfMemoryError(t)) {
        String message = "An OutOfMemoryException occurred. Consider specifying higher memory " + "instances in PipelineOptions.\n";
        LOG.error("{}: {}", logPrefix, message);
        error.setMessage(message + DataflowWorkerLoggingHandler.formatException(t));
    } else {
        LOG.error("{}: Uncaught exception occurred during work unit execution. This will be retried.", logPrefix, t);
        error.setMessage(DataflowWorkerLoggingHandler.formatException(t));
    }
    status.setErrors(ImmutableList.of(error));
    return execute(status);
}

Example 3 with UserCodeException

use of org.apache.beam.sdk.util.UserCodeException in project beam by apache.

the class ApexParDoOperator method processElementInReadyWindows.

private Iterable<WindowedValue<InputT>> processElementInReadyWindows(WindowedValue<InputT> elem) {
    try {
        pushbackDoFnRunner.startBundle();
        if (currentKeyStateInternals != null) {
            InputT value = elem.getValue();
            final Object key;
            final Coder<Object> keyCoder;
            @SuppressWarnings({ "rawtypes", "unchecked" }) WindowedValueCoder<InputT> wvCoder = (WindowedValueCoder) inputCoder;
            if (value instanceof KeyedWorkItem) {
                key = ((KeyedWorkItem) value).key();
                @SuppressWarnings({ "rawtypes", "unchecked" }) KeyedWorkItemCoder<Object, ?> kwiCoder = (KeyedWorkItemCoder) wvCoder.getValueCoder();
                keyCoder = kwiCoder.getKeyCoder();
            } else {
                key = ((KV) value).getKey();
                @SuppressWarnings({ "rawtypes", "unchecked" }) KvCoder<Object, ?> kwiCoder = (KvCoder) wvCoder.getValueCoder();
                keyCoder = kwiCoder.getKeyCoder();
            }
            ((StateInternalsProxy) currentKeyStateInternals).setKey(key);
            currentKeyTimerInternals.setContext(key, keyCoder, new Instant(this.currentInputWatermark), new Instant(this.currentOutputWatermark));
        }
        Iterable<WindowedValue<InputT>> pushedBack = pushbackDoFnRunner.processElementInReadyWindows(elem);
        pushbackDoFnRunner.finishBundle();
        return pushedBack;
    } catch (UserCodeException ue) {
        if (ue.getCause() instanceof AssertionError) {
            ApexRunner.ASSERTION_ERROR.set((AssertionError) ue.getCause());
        }
        throw ue;
    }
}

Example 4 with UserCodeException

use of org.apache.beam.sdk.util.UserCodeException in project beam by apache.

the class DirectRunner method run.

@Override
public DirectPipelineResult run(Pipeline pipeline) {
    try {
        options = MAPPER.readValue(MAPPER.writeValueAsBytes(options), PipelineOptions.class).as(DirectOptions.class);
    } catch (IOException e) {
        throw new IllegalArgumentException("PipelineOptions specified failed to serialize to JSON.", e);
    }
    performRewrites(pipeline);
    MetricsEnvironment.setMetricsSupported(true);
    try {
        DirectGraphVisitor graphVisitor = new DirectGraphVisitor();
        pipeline.traverseTopologically(graphVisitor);
        @SuppressWarnings("rawtypes") KeyedPValueTrackingVisitor keyedPValueVisitor = KeyedPValueTrackingVisitor.create();
        pipeline.traverseTopologically(keyedPValueVisitor);
        DisplayDataValidator.validatePipeline(pipeline);
        DisplayDataValidator.validateOptions(options);
        ExecutorService metricsPool = Executors.newCachedThreadPool(new ThreadFactoryBuilder().setThreadFactory(MoreExecutors.platformThreadFactory()).setDaemon(// otherwise you say you want to leak, please don't!
        false).setNameFormat("direct-metrics-counter-committer").build());
        DirectGraph graph = graphVisitor.getGraph();
        EvaluationContext context = EvaluationContext.create(clockSupplier.get(), Enforcement.bundleFactoryFor(enabledEnforcements, graph), graph, keyedPValueVisitor.getKeyedPValues(), metricsPool);
        TransformEvaluatorRegistry registry = TransformEvaluatorRegistry.javaSdkNativeRegistry(context, options);
        PipelineExecutor executor = ExecutorServiceParallelExecutor.create(options.getTargetParallelism(), registry, Enforcement.defaultModelEnforcements(enabledEnforcements), context, metricsPool);
        executor.start(graph, RootProviderRegistry.javaNativeRegistry(context, options));
        DirectPipelineResult result = new DirectPipelineResult(executor, context);
        if (options.isBlockOnRun()) {
            try {
                result.waitUntilFinish();
            } catch (UserCodeException userException) {
                throw new PipelineExecutionException(userException.getCause());
            } catch (Throwable t) {
                if (t instanceof RuntimeException) {
                    throw (RuntimeException) t;
                }
                throw new RuntimeException(t);
            }
        }
        return result;
    } finally {
        MetricsEnvironment.setMetricsSupported(false);
    }
}

Example 5 with UserCodeException

use of org.apache.beam.sdk.util.UserCodeException 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);
    }
}