Examples with DriverState org.apache.beam.runners.local.ExecutionDriver.DriverState used on opensource projects

Example 1 with DriverState

use of org.apache.beam.runners.local.ExecutionDriver.DriverState in project beam by apache.

the class ExecutorServiceParallelExecutor method start.

@Override
// TODO: [BEAM-4563] Pass Future back to consumer to check for async errors
@SuppressWarnings("FutureReturnValueIgnored")
public void start(DirectGraph graph, RootProviderRegistry rootProviderRegistry) {
    int numTargetSplits = Math.max(3, targetParallelism);
    ImmutableMap.Builder<AppliedPTransform<?, ?, ?>, Queue<CommittedBundle<?>>> pendingRootBundles = ImmutableMap.builder();
    for (AppliedPTransform<?, ?, ?> root : graph.getRootTransforms()) {
        Queue<CommittedBundle<?>> pending = Queues.newArrayDeque();
        try {
            Collection<CommittedBundle<?>> initialInputs = rootProviderRegistry.getInitialInputs(root, numTargetSplits);
            pending.addAll(initialInputs);
        } catch (Exception e) {
            throw UserCodeException.wrap(e);
        }
        pendingRootBundles.put(root, pending);
    }
    evaluationContext.initialize(pendingRootBundles.build());
    final ExecutionDriver executionDriver = QuiescenceDriver.create(evaluationContext, graph, this, visibleUpdates, pendingRootBundles.build());
    executorService.submit(new Runnable() {

        @Override
        public void run() {
            DriverState drive = executionDriver.drive();
            if (drive.isTermainal()) {
                State newPipelineState = State.UNKNOWN;
                switch(drive) {
                    case FAILED:
                        newPipelineState = State.FAILED;
                        break;
                    case SHUTDOWN:
                        newPipelineState = State.DONE;
                        break;
                    case CONTINUE:
                        throw new IllegalStateException(String.format("%s should not be a terminal state", DriverState.CONTINUE));
                    default:
                        throw new IllegalArgumentException(String.format("Unknown %s %s", DriverState.class.getSimpleName(), drive));
                }
                shutdownIfNecessary(newPipelineState);
            } else {
                executorService.submit(this);
            }
        }
    });
}