Examples with OutputStatisticsReporter org.apache.tez.runtime.api.OutputStatisticsReporter used on opensource projects

Example 1 with OutputStatisticsReporter

use of org.apache.tez.runtime.api.OutputStatisticsReporter in project tez by apache.

the class TestOnFileSortedOutput method createTezOutputContext.

private OutputContext createTezOutputContext() throws IOException {
    String[] workingDirs = { workingDir.toString() };
    UserPayload payLoad = TezUtils.createUserPayloadFromConf(conf);
    DataOutputBuffer serviceProviderMetaData = new DataOutputBuffer();
    serviceProviderMetaData.writeInt(PORT);
    TezCounters counters = new TezCounters();
    OutputStatisticsReporter reporter = mock(OutputStatisticsReporter.class);
    doAnswer(new Answer() {

        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            outputSize.set((Long) invocation.getArguments()[0]);
            return null;
        }
    }).when(reporter).reportDataSize(anyLong());
    doAnswer(new Answer() {

        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            numRecords.set((Long) invocation.getArguments()[0]);
            return null;
        }
    }).when(reporter).reportItemsProcessed(anyLong());
    OutputContext context = mock(OutputContext.class);
    doReturn(counters).when(context).getCounters();
    doReturn(workingDirs).when(context).getWorkDirs();
    doReturn(payLoad).when(context).getUserPayload();
    doReturn(5 * 1024 * 1024l).when(context).getTotalMemoryAvailableToTask();
    doReturn(UniqueID).when(context).getUniqueIdentifier();
    doReturn("v1").when(context).getDestinationVertexName();
    doReturn(ByteBuffer.wrap(serviceProviderMetaData.getData())).when(context).getServiceProviderMetaData(conf.get(TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID, TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID_DEFAULT));
    doAnswer(new Answer() {

        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            long requestedSize = (Long) invocation.getArguments()[0];
            MemoryUpdateCallbackHandler callback = (MemoryUpdateCallbackHandler) invocation.getArguments()[1];
            callback.memoryAssigned(requestedSize);
            return null;
        }
    }).when(context).requestInitialMemory(anyLong(), any(MemoryUpdateCallback.class));
    ExecutionContext ExecutionContext = mock(ExecutionContext.class);
    doReturn(HOST).when(ExecutionContext).getHostName();
    doReturn(reporter).when(context).getStatisticsReporter();
    doReturn(ExecutionContext).when(context).getExecutionContext();
    return context;
}

Example 2 with OutputStatisticsReporter

use of org.apache.tez.runtime.api.OutputStatisticsReporter in project tez by apache.

the class OutputTestHelpers method createOutputContext.

static OutputContext createOutputContext(Configuration conf, Path workingDir) throws IOException {
    OutputContext ctx = mock(OutputContext.class);
    doAnswer(new Answer<Void>() {

        @Override
        public Void answer(InvocationOnMock invocation) throws Throwable {
            long requestedSize = (Long) invocation.getArguments()[0];
            MemoryUpdateCallbackHandler callback = (MemoryUpdateCallbackHandler) invocation.getArguments()[1];
            callback.memoryAssigned(requestedSize);
            return null;
        }
    }).when(ctx).requestInitialMemory(anyLong(), any(MemoryUpdateCallback.class));
    doReturn(TezUtils.createUserPayloadFromConf(conf)).when(ctx).getUserPayload();
    doReturn("destinationVertex").when(ctx).getDestinationVertexName();
    doReturn("UUID").when(ctx).getUniqueIdentifier();
    doReturn(new String[] { workingDir.toString() }).when(ctx).getWorkDirs();
    doReturn(200 * 1024 * 1024l).when(ctx).getTotalMemoryAvailableToTask();
    doReturn(new TezCounters()).when(ctx).getCounters();
    OutputStatisticsReporter statsReporter = mock(OutputStatisticsReporter.class);
    doReturn(statsReporter).when(ctx).getStatisticsReporter();
    doReturn(new ExecutionContextImpl("localhost")).when(ctx).getExecutionContext();
    return ctx;
}

Example 3 with OutputStatisticsReporter

use of org.apache.tez.runtime.api.OutputStatisticsReporter in project tez by apache.

the class OutputTestHelpers method createOutputContext.

static OutputContext createOutputContext() throws IOException {
    OutputContext outputContext = mock(OutputContext.class);
    Configuration conf = new TezConfiguration();
    UserPayload payLoad = TezUtils.createUserPayloadFromConf(conf);
    String[] workingDirs = new String[] { "workDir1" };
    OutputStatisticsReporter statsReporter = mock(OutputStatisticsReporter.class);
    TezCounters counters = new TezCounters();
    doReturn("destinationVertex").when(outputContext).getDestinationVertexName();
    doReturn(payLoad).when(outputContext).getUserPayload();
    doReturn(workingDirs).when(outputContext).getWorkDirs();
    doReturn(200 * 1024 * 1024l).when(outputContext).getTotalMemoryAvailableToTask();
    doReturn(counters).when(outputContext).getCounters();
    doReturn(statsReporter).when(outputContext).getStatisticsReporter();
    return outputContext;
}