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

Example 6 with MemoryUpdateCallbackHandler

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

the class TestDefaultSorter method basicTest.

@Test(timeout = 30000)
public // Test TEZ-1977
void basicTest() throws IOException {
    OutputContext context = createTezOutputContext();
    MemoryUpdateCallbackHandler handler = new MemoryUpdateCallbackHandler();
    try {
        // Setting IO_SORT_MB to greater than available mem limit (should throw exception)
        conf.setInt(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB, 300);
        context.requestInitialMemory(ExternalSorter.getInitialMemoryRequirement(conf, context.getTotalMemoryAvailableToTask()), new MemoryUpdateCallbackHandler());
        fail();
    } catch (IllegalArgumentException e) {
        assertTrue(e.getMessage().contains(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB));
    }
    conf.setLong(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB, 1);
    context.requestInitialMemory(ExternalSorter.getInitialMemoryRequirement(conf, context.getTotalMemoryAvailableToTask()), handler);
    DefaultSorter sorter = new DefaultSorter(context, conf, 5, handler.getMemoryAssigned());
    // Write 1000 keys each of size 1000, (> 1 spill should happen)
    try {
        writeData(sorter, 1000, 1000);
        assertTrue(sorter.getNumSpills() > 2);
        verifyCounters(sorter, context);
    } catch (IOException ioe) {
        fail(ioe.getMessage());
    }
    verifyOutputPermissions(context.getUniqueIdentifier());
}

Example 7 with MemoryUpdateCallbackHandler

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

the class TestOnFileUnorderedKVOutput method createOutputContext.

private OutputContext createOutputContext(Configuration conf, TezSharedExecutor sharedExecutor) throws IOException {
    int appAttemptNumber = 1;
    TezUmbilical tezUmbilical = mock(TezUmbilical.class);
    String dagName = "currentDAG";
    String taskVertexName = "currentVertex";
    String destinationVertexName = "destinationVertex";
    TezDAGID dagID = TezDAGID.getInstance("2000", 1, 1);
    TezVertexID vertexID = TezVertexID.getInstance(dagID, 1);
    TezTaskID taskID = TezTaskID.getInstance(vertexID, 1);
    TezTaskAttemptID taskAttemptID = TezTaskAttemptID.getInstance(taskID, 1);
    UserPayload userPayload = TezUtils.createUserPayloadFromConf(conf);
    TaskSpec mockSpec = mock(TaskSpec.class);
    when(mockSpec.getInputs()).thenReturn(Collections.singletonList(mock(InputSpec.class)));
    when(mockSpec.getOutputs()).thenReturn(Collections.singletonList(mock(OutputSpec.class)));
    task = new LogicalIOProcessorRuntimeTask(mockSpec, appAttemptNumber, new Configuration(), new String[] { "/" }, tezUmbilical, null, null, null, null, "", null, 1024, false, new DefaultHadoopShim(), sharedExecutor);
    LogicalIOProcessorRuntimeTask runtimeTask = spy(task);
    Map<String, String> auxEnv = new HashMap<String, String>();
    ByteBuffer bb = ByteBuffer.allocate(4);
    bb.putInt(shufflePort);
    bb.position(0);
    AuxiliaryServiceHelper.setServiceDataIntoEnv(conf.get(TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID, TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID_DEFAULT), bb, auxEnv);
    OutputDescriptor outputDescriptor = mock(OutputDescriptor.class);
    when(outputDescriptor.getClassName()).thenReturn("OutputDescriptor");
    OutputContext realOutputContext = new TezOutputContextImpl(conf, new String[] { workDir.toString() }, appAttemptNumber, tezUmbilical, dagName, taskVertexName, destinationVertexName, -1, taskAttemptID, 0, userPayload, runtimeTask, null, auxEnv, new MemoryDistributor(1, 1, conf), outputDescriptor, null, new ExecutionContextImpl("localhost"), 2048, new TezSharedExecutor(defaultConf));
    verify(runtimeTask, times(1)).addAndGetTezCounter(destinationVertexName);
    verify(runtimeTask, times(1)).getTaskStatistics();
    // verify output stats object got created
    Assert.assertTrue(task.getTaskStatistics().getIOStatistics().containsKey(destinationVertexName));
    OutputContext outputContext = spy(realOutputContext);
    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(outputContext).requestInitialMemory(anyLong(), any(MemoryUpdateCallback.class));
    return outputContext;
}

Example 8 with MemoryUpdateCallbackHandler

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

the class TestOrderedGroupedKVInput method createMockInputContext.

private InputContext createMockInputContext() throws IOException {
    InputContext inputContext = mock(InputContext.class);
    Configuration conf = new TezConfiguration();
    UserPayload payLoad = TezUtils.createUserPayloadFromConf(conf);
    String[] workingDirs = new String[] { "workDir1" };
    TezCounters counters = new TezCounters();
    doReturn(payLoad).when(inputContext).getUserPayload();
    doReturn(workingDirs).when(inputContext).getWorkDirs();
    doReturn(200 * 1024 * 1024l).when(inputContext).getTotalMemoryAvailableToTask();
    doReturn(counters).when(inputContext).getCounters();
    doAnswer(new Answer() {

        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            Object[] args = invocation.getArguments();
            if (args[1] instanceof MemoryUpdateCallbackHandler) {
                MemoryUpdateCallbackHandler memUpdateCallbackHandler = (MemoryUpdateCallbackHandler) args[1];
                memUpdateCallbackHandler.memoryAssigned(200 * 1024 * 1024);
            } else {
                Assert.fail();
            }
            return null;
        }
    }).when(inputContext).requestInitialMemory(any(long.class), any(MemoryUpdateCallbackHandler.class));
    return inputContext;
}

Example 9 with MemoryUpdateCallbackHandler

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

the class OrderedGroupedKVInput method initialize.

@Override
public synchronized List<Event> initialize() throws IOException {
    this.conf = TezUtils.createConfFromUserPayload(getContext().getUserPayload());
    if (this.getNumPhysicalInputs() == 0) {
        getContext().requestInitialMemory(0l, null);
        isStarted.set(true);
        getContext().inputIsReady();
        LOG.info("input fetch not required since there are 0 physical inputs for input vertex: " + getContext().getSourceVertexName());
        return Collections.emptyList();
    }
    long initialMemoryRequest = Shuffle.getInitialMemoryRequirement(conf, getContext().getTotalMemoryAvailableToTask());
    this.memoryUpdateCallbackHandler = new MemoryUpdateCallbackHandler();
    getContext().requestInitialMemory(initialMemoryRequest, memoryUpdateCallbackHandler);
    this.inputKeyCounter = getContext().getCounters().findCounter(TaskCounter.REDUCE_INPUT_GROUPS);
    this.inputValueCounter = getContext().getCounters().findCounter(TaskCounter.REDUCE_INPUT_RECORDS);
    this.shuffledInputs = getContext().getCounters().findCounter(TaskCounter.NUM_SHUFFLED_INPUTS);
    this.conf.setStrings(TezRuntimeFrameworkConfigs.LOCAL_DIRS, getContext().getWorkDirs());
    return Collections.emptyList();
}

Example 10 with MemoryUpdateCallbackHandler

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

the class UnorderedKVInput method initialize.

@Override
public synchronized List<Event> initialize() throws Exception {
    Preconditions.checkArgument(getNumPhysicalInputs() != -1, "Number of Inputs has not been set");
    this.conf = TezUtils.createConfFromUserPayload(getContext().getUserPayload());
    if (getNumPhysicalInputs() == 0) {
        getContext().requestInitialMemory(0l, null);
        isStarted.set(true);
        getContext().inputIsReady();
        LOG.info("input fetch not required since there are 0 physical inputs for input vertex: " + getContext().getSourceVertexName());
        return Collections.emptyList();
    } else {
        long initalMemReq = getInitialMemoryReq();
        memoryUpdateCallbackHandler = new MemoryUpdateCallbackHandler();
        this.getContext().requestInitialMemory(initalMemReq, memoryUpdateCallbackHandler);
    }
    this.conf.setStrings(TezRuntimeFrameworkConfigs.LOCAL_DIRS, getContext().getWorkDirs());
    this.inputRecordCounter = getContext().getCounters().findCounter(TaskCounter.INPUT_RECORDS_PROCESSED);
    return Collections.emptyList();
}