Examples with MemoryUpdateCallbackHandler org.apache.tez.runtime.library.common.MemoryUpdateCallbackHandler used on opensource projects

Example 11 with MemoryUpdateCallbackHandler

use of org.apache.tez.runtime.library.common.MemoryUpdateCallbackHandler in project tez by apache.

the class UnorderedPartitionedKVOutput method initialize.

@Override
public synchronized List<Event> initialize() throws Exception {
    this.conf = TezUtils.createConfFromUserPayload(getContext().getUserPayload());
    this.conf.setStrings(TezRuntimeFrameworkConfigs.LOCAL_DIRS, getContext().getWorkDirs());
    this.conf.setInt(TezRuntimeFrameworkConfigs.TEZ_RUNTIME_NUM_EXPECTED_PARTITIONS, getNumPhysicalOutputs());
    this.memoryUpdateCallbackHandler = new MemoryUpdateCallbackHandler();
    getContext().requestInitialMemory(UnorderedPartitionedKVWriter.getInitialMemoryRequirement(conf, getContext().getTotalMemoryAvailableToTask()), memoryUpdateCallbackHandler);
    return Collections.emptyList();
}

Example 12 with MemoryUpdateCallbackHandler

use of org.apache.tez.runtime.library.common.MemoryUpdateCallbackHandler in project tez by apache.

the class UnorderedKVOutput method initialize.

@Override
public synchronized List<Event> initialize() throws Exception {
    this.conf = TezUtils.createConfFromUserPayload(getContext().getUserPayload());
    this.conf.setStrings(TezRuntimeFrameworkConfigs.LOCAL_DIRS, getContext().getWorkDirs());
    this.conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_PARTITIONER_CLASS, CustomPartitioner.class.getName());
    this.memoryUpdateCallbackHandler = new MemoryUpdateCallbackHandler();
    boolean pipelinedShuffle = this.conf.getBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_PIPELINED_SHUFFLE_ENABLED, TezRuntimeConfiguration.TEZ_RUNTIME_PIPELINED_SHUFFLE_ENABLED_DEFAULT);
    long memRequestSize = (pipelinedShuffle) ? UnorderedPartitionedKVWriter.getInitialMemoryRequirement(conf, getContext().getTotalMemoryAvailableToTask()) : 0;
    getContext().requestInitialMemory(memRequestSize, memoryUpdateCallbackHandler);
    return Collections.emptyList();
}

Example 13 with MemoryUpdateCallbackHandler

use of org.apache.tez.runtime.library.common.MemoryUpdateCallbackHandler in project tez by apache.

the class TestDefaultSorter method testWithMultipleSpillsWithFinalMergeDisabled.

@Test(timeout = 60000)
@SuppressWarnings("unchecked")
public void testWithMultipleSpillsWithFinalMergeDisabled() throws IOException {
    OutputContext context = createTezOutputContext();
    conf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_ENABLE_FINAL_MERGE_IN_OUTPUT, false);
    conf.setLong(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB, 4);
    conf.setInt(TezRuntimeConfiguration.TEZ_RUNTIME_INDEX_CACHE_MEMORY_LIMIT_BYTES, 1);
    MemoryUpdateCallbackHandler handler = new MemoryUpdateCallbackHandler();
    context.requestInitialMemory(ExternalSorter.getInitialMemoryRequirement(conf, context.getTotalMemoryAvailableToTask()), handler);
    DefaultSorter sorter = new DefaultSorter(context, conf, 1, handler.getMemoryAssigned());
    writeData(sorter, 10000, 1000);
    int spillCount = sorter.getNumSpills();
    ArgumentCaptor<List> eventCaptor = ArgumentCaptor.forClass(List.class);
    verify(context, times(1)).sendEvents(eventCaptor.capture());
    List<Event> events = eventCaptor.getValue();
    int spillIndex = 0;
    for (Event event : events) {
        if (event instanceof CompositeDataMovementEvent) {
            CompositeDataMovementEvent cdme = (CompositeDataMovementEvent) event;
            ShuffleUserPayloads.DataMovementEventPayloadProto shufflePayload = ShuffleUserPayloads.DataMovementEventPayloadProto.parseFrom(ByteString.copyFrom(cdme.getUserPayload()));
            assertTrue(shufflePayload.getPathComponent().equalsIgnoreCase(UniqueID + "_" + spillIndex));
            verifyOutputPermissions(shufflePayload.getPathComponent());
            spillIndex++;
        }
    }
    assertTrue(spillIndex == spillCount);
    verifyCounters(sorter, context);
}

Example 14 with MemoryUpdateCallbackHandler

use of org.apache.tez.runtime.library.common.MemoryUpdateCallbackHandler in project tez by apache.

the class TestDefaultSorter 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();
    OutputContext context = mock(OutputContext.class);
    ExecutionContext execContext = new ExecutionContextImpl("localhost");
    doReturn(mock(OutputStatisticsReporter.class)).when(context).getStatisticsReporter();
    doReturn(execContext).when(context).getExecutionContext();
    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));
    return context;
}

Example 15 with MemoryUpdateCallbackHandler

use of org.apache.tez.runtime.library.common.MemoryUpdateCallbackHandler in project tez by apache.

the class TestDefaultSorter method testPartitionStats.

void testPartitionStats(boolean withStats) throws IOException {
    conf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_REPORT_PARTITION_STATS, withStats);
    OutputContext context = createTezOutputContext();
    conf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_ENABLE_FINAL_MERGE_IN_OUTPUT, false);
    conf.setLong(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB, 4);
    MemoryUpdateCallbackHandler handler = new MemoryUpdateCallbackHandler();
    context.requestInitialMemory(ExternalSorter.getInitialMemoryRequirement(conf, context.getTotalMemoryAvailableToTask()), handler);
    DefaultSorter sorter = new DefaultSorter(context, conf, 1, handler.getMemoryAssigned());
    writeData(sorter, 1000, 10);
    assertTrue(sorter.getNumSpills() == 1);
    verifyCounters(sorter, context);
    if (withStats) {
        assertTrue(sorter.getPartitionStats() != null);
    } else {
        assertTrue(sorter.getPartitionStats() == null);
    }
}