Examples with NameContext org.apache.beam.runners.dataflow.worker.counters.NameContext used on opensource projects

Example 1 with NameContext

use of org.apache.beam.runners.dataflow.worker.counters.NameContext in project beam by apache.

the class DataflowElementExecutionTrackerTest method testTypicalUsage.

/**
 * Typical usage scenario.
 */
@Test
public void testTypicalUsage() throws IOException {
    NameContext stepA = createStep("A");
    NameContext stepB = createStep("B");
    NameContext stepC = createStep("C");
    NameContext stepD = createStep("D");
    // Comments track journal of executions for next sample, and partial timings not yet reported
    // IDLE A1 | {}
    tracker.enter(stepA);
    // IDLE A1 B1 | {}
    tracker.enter(stepB);
    // IDLE A1 B1 A1 | {}
    tracker.exit();
    // A1 | {A1:2}
    tracker.takeSample(40);
    assertThat(getCounterValue(stepB), equalTo(distribution(10)));
    // A1 B2 | {A1:2}
    tracker.enter(stepB);
    // A1 B2 A1 | {A1:2}
    tracker.exit();
    // A1 B2 A1 C1 | {A1:2}
    tracker.enter(stepC);
    // A1 B2 A1 C1 D1 | {A1:2}
    tracker.enter(stepD);
    // D1 | {A1:4 C1:1 D1:1}
    tracker.takeSample(50);
    assertThat(getCounterValue(stepB), equalTo(distribution(10, 10)));
    // D1 C1 | {A1:4 C1:1 D1:1}
    tracker.exit();
    // D1 C1 A1 | {A1:4 C1:1 D1:1}
    tracker.exit();
    // D1 C1 A1 C2 | {A1:4 C1:1 D1:1}
    tracker.enter(stepC);
    // C2 | {A1:5 C2:1}
    tracker.takeSample(40);
    assertThat(getCounterValue(stepC), equalTo(distribution(20)));
    assertThat(getCounterValue(stepD), equalTo(distribution(20)));
    // C2 A1 | {A1:5 C2:1}
    tracker.exit();
    // C2 A1 IDLE | {A1:5 C2:1}
    tracker.exit();
    // done
    tracker.takeSample(30);
    assertThat(getCounterValue(stepA), equalTo(distribution(60)));
    assertThat(getCounterValue(stepB), equalTo(distribution(10, 10)));
    assertThat(getCounterValue(stepC), equalTo(distribution(20, 20)));
    assertThat(getCounterValue(stepD), equalTo(distribution(20)));
}

Example 2 with NameContext

use of org.apache.beam.runners.dataflow.worker.counters.NameContext in project beam by apache.

the class DataflowElementExecutionTrackerTest method testDisabledByExperiment.

/**
 * Ensure functionality is correctly disabled when the experiment is not set.
 */
@Test
public void testDisabledByExperiment() throws IOException {
    List<String> experiments = options.getExperiments();
    experiments.remove(DataflowElementExecutionTracker.TIME_PER_ELEMENT_EXPERIMENT);
    options.setExperiments(experiments);
    tracker = DataflowElementExecutionTracker.create(counters, options);
    NameContext step = createStep("A");
    tracker.enter(step);
    tracker.exit();
    tracker.takeSample(10);
    assertThat(getCounter(step), nullValue());
}

Example 3 with NameContext

use of org.apache.beam.runners.dataflow.worker.counters.NameContext in project beam by apache.

the class DataflowElementExecutionTrackerTest method testElementsTrackedIndividuallyForAStep.

/**
 * Verify that each entry into a step tracks a separate element execution.
 */
@Test
public void testElementsTrackedIndividuallyForAStep() throws IOException {
    NameContext stepA = createStep("A");
    NameContext stepB = createStep("B");
    tracker.enter(stepA);
    tracker.enter(stepB);
    tracker.exit();
    tracker.enter(stepB);
    tracker.exit();
    tracker.exit();
    // Expected journal: IDLE A1 B1 A1 B2 A1 IDLE
    tracker.takeSample(70);
    assertThat(getCounterValue(stepA), equalTo(distribution(30)));
    assertThat(getCounterValue(stepB), equalTo(distribution(10, 10)));
}

Example 4 with NameContext

use of org.apache.beam.runners.dataflow.worker.counters.NameContext in project beam by apache.

the class BeamFnMapTaskExecutorFactory method createOperationTransformForExecutableStageNode.

private Function<Node, Node> createOperationTransformForExecutableStageNode(final Network<Node, Edge> network, final String stageName, final DataflowExecutionContext<?> executionContext, final JobBundleFactory jobBundleFactory) {
    return new TypeSafeNodeFunction<ExecutableStageNode>(ExecutableStageNode.class) {

        @Override
        public Node typedApply(ExecutableStageNode input) {
            StageBundleFactory stageBundleFactory = jobBundleFactory.forStage(input.getExecutableStage());
            Iterable<OutputReceiverNode> outputReceiverNodes = Iterables.filter(network.successors(input), OutputReceiverNode.class);
            Map<String, OutputReceiver> outputReceiverMap = new HashMap<>();
            Lists.newArrayList(outputReceiverNodes).stream().forEach(outputReceiverNode -> outputReceiverMap.put(outputReceiverNode.getPcollectionId(), outputReceiverNode.getOutputReceiver()));
            ImmutableMap.Builder<String, DataflowOperationContext> ptransformIdToOperationContextBuilder = ImmutableMap.builder();
            for (Map.Entry<String, NameContext> entry : input.getPTransformIdToPartialNameContextMap().entrySet()) {
                NameContext fullNameContext = NameContext.create(stageName, entry.getValue().originalName(), entry.getValue().systemName(), entry.getValue().userName());
                DataflowOperationContext operationContext = executionContext.createOperationContext(fullNameContext);
                ptransformIdToOperationContextBuilder.put(entry.getKey(), operationContext);
            }
            ImmutableMap<String, DataflowOperationContext> ptransformIdToOperationContexts = ptransformIdToOperationContextBuilder.build();
            ImmutableMap<String, SideInputReader> ptransformIdToSideInputReaders = buildPTransformIdToSideInputReadersMap(executionContext, input, ptransformIdToOperationContexts);
            Map<RunnerApi.ExecutableStagePayload.SideInputId, PCollectionView<?>> ptransformIdToSideInputIdToPCollectionView = buildSideInputIdToPCollectionView(input);
            return OperationNode.create(new ProcessRemoteBundleOperation(input.getExecutableStage(), executionContext.createOperationContext(NameContext.create(stageName, stageName, stageName, stageName)), stageBundleFactory, outputReceiverMap, ptransformIdToSideInputReaders, ptransformIdToSideInputIdToPCollectionView));
        }
    };
}

Example 5 with NameContext

use of org.apache.beam.runners.dataflow.worker.counters.NameContext in project beam by apache.

the class BeamFnMapTaskExecutorFactory method createOperationTransformForParallelInstructionNodes.

/**
 * Creates an {@link Operation} from the given {@link ParallelInstruction} definition using the
 * provided {@link ReaderFactory}.
 */
Function<Node, Node> createOperationTransformForParallelInstructionNodes(final String stageName, final Network<Node, Edge> network, final PipelineOptions options, final ReaderFactory readerFactory, final SinkFactory sinkFactory, final DataflowExecutionContext<?> executionContext) {
    return new TypeSafeNodeFunction<ParallelInstructionNode>(ParallelInstructionNode.class) {

        @Override
        public Node typedApply(ParallelInstructionNode node) {
            ParallelInstruction instruction = node.getParallelInstruction();
            NameContext nameContext = NameContext.create(stageName, instruction.getOriginalName(), instruction.getSystemName(), instruction.getName());
            try {
                DataflowOperationContext context = executionContext.createOperationContext(nameContext);
                if (instruction.getRead() != null) {
                    return createReadOperation(network, node, options, readerFactory, executionContext, context);
                } else if (instruction.getWrite() != null) {
                    return createWriteOperation(node, options, sinkFactory, executionContext, context);
                } else if (instruction.getParDo() != null) {
                    return createParDoOperation(network, node, options, executionContext, context);
                } else if (instruction.getPartialGroupByKey() != null) {
                    return createPartialGroupByKeyOperation(network, node, options, executionContext, context);
                } else if (instruction.getFlatten() != null) {
                    return createFlattenOperation(network, node, context);
                } else {
                    throw new IllegalArgumentException(String.format("Unexpected instruction: %s", instruction));
                }
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }
    };
}