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

Example 6 with FetchedInputAllocator

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

the class TestShuffleInputEventHandlerImpl method createShuffleManager.

private ShuffleManager createShuffleManager(InputContext inputContext) throws IOException {
    Path outDirBase = new Path(".", "outDir");
    String[] outDirs = new String[] { outDirBase.toString() };
    doReturn(outDirs).when(inputContext).getWorkDirs();
    conf.setStrings(TezRuntimeFrameworkConfigs.LOCAL_DIRS, inputContext.getWorkDirs());
    DataOutputBuffer out = new DataOutputBuffer();
    Token<JobTokenIdentifier> token = new Token<JobTokenIdentifier>(new JobTokenIdentifier(), new JobTokenSecretManager(null));
    token.write(out);
    doReturn(ByteBuffer.wrap(out.getData())).when(inputContext).getServiceConsumerMetaData(conf.get(TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID, TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID_DEFAULT));
    FetchedInputAllocator inputAllocator = mock(FetchedInputAllocator.class);
    ShuffleManager realShuffleManager = new ShuffleManager(inputContext, conf, 2, 1024, false, -1, null, inputAllocator);
    ShuffleManager shuffleManager = spy(realShuffleManager);
    return shuffleManager;
}

Example 7 with FetchedInputAllocator

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

the class TestShuffleInputEventHandlerImpl method testMultipleEvents1.

@Test(timeout = 5000)
public void testMultipleEvents1() throws IOException {
    InputContext inputContext = mock(InputContext.class);
    ShuffleManager shuffleManager = mock(ShuffleManager.class);
    FetchedInputAllocator inputAllocator = mock(FetchedInputAllocator.class);
    ShuffleInputEventHandlerImpl handler = new ShuffleInputEventHandlerImpl(inputContext, shuffleManager, inputAllocator, null, false, 0, false);
    int taskIndex1 = 1;
    Event dme1 = createDataMovementEvent(0, taskIndex1, createEmptyPartitionByteString(0));
    int taskIndex2 = 2;
    Event dme2 = createDataMovementEvent(0, taskIndex2, null);
    List<Event> eventList = new LinkedList<Event>();
    eventList.add(dme1);
    eventList.add(dme2);
    handler.handleEvents(eventList);
    InputAttemptIdentifier expectedIdentifier1 = new InputAttemptIdentifier(taskIndex1, 0);
    CompositeInputAttemptIdentifier expectedIdentifier2 = new CompositeInputAttemptIdentifier(taskIndex2, 0, PATH_COMPONENT, 1);
    verify(shuffleManager).addCompletedInputWithNoData(eq(expectedIdentifier1));
    verify(shuffleManager).addKnownInput(eq(HOST), eq(PORT), eq(expectedIdentifier2), eq(0));
}

Example 8 with FetchedInputAllocator

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

the class TestShuffleInputEventHandlerImpl method testOtherPartitionEmpty.

@Test(timeout = 5000)
public void testOtherPartitionEmpty() throws IOException {
    InputContext inputContext = mock(InputContext.class);
    ShuffleManager shuffleManager = mock(ShuffleManager.class);
    FetchedInputAllocator inputAllocator = mock(FetchedInputAllocator.class);
    ShuffleInputEventHandlerImpl handler = new ShuffleInputEventHandlerImpl(inputContext, shuffleManager, inputAllocator, null, false, 0, false);
    int taskIndex = 1;
    Event dme = createDataMovementEvent(0, taskIndex, createEmptyPartitionByteString(1));
    List<Event> eventList = new LinkedList<Event>();
    eventList.add(dme);
    handler.handleEvents(eventList);
    CompositeInputAttemptIdentifier expectedIdentifier = new CompositeInputAttemptIdentifier(taskIndex, 0, PATH_COMPONENT, 1);
    verify(shuffleManager).addKnownInput(eq(HOST), eq(PORT), eq(expectedIdentifier), eq(0));
}

Example 9 with FetchedInputAllocator

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

the class TestShuffleManager method testMultiplePartitions.

/**
 * One reducer fetches multiple partitions from each mapper.
 * For a given mapper, the reducer sends DataMovementEvents for several
 * partitions, wait for some time and then send DataMovementEvents for the
 * rest of the partitions. Then do the same thing for the next mapper.
 * Verify ShuffleManager is able to get all the events.
 */
@Test(timeout = 50000)
public void testMultiplePartitions() throws Exception {
    final int numOfMappers = 3;
    final int numOfPartitions = 5;
    final int firstPart = 2;
    InputContext inputContext = createInputContext();
    ShuffleManagerForTest shuffleManager = createShuffleManager(inputContext, numOfMappers * numOfPartitions);
    FetchedInputAllocator inputAllocator = mock(FetchedInputAllocator.class);
    ShuffleInputEventHandlerImpl handler = new ShuffleInputEventHandlerImpl(inputContext, shuffleManager, inputAllocator, null, false, 0, false);
    shuffleManager.run();
    List<Event> eventList = new LinkedList<Event>();
    // The physical input index within the reduce task
    int targetIndex = 0;
    for (int i = 0; i < numOfMappers; i++) {
        String mapperHost = "host" + i;
        // The physical output index within the map task
        int srcIndex = 20;
        // Send the first batch of DataMovementEvents
        eventList.clear();
        for (int j = 0; j < firstPart; j++) {
            Event dme = createDataMovementEvent(mapperHost, srcIndex++, targetIndex++);
            eventList.add(dme);
        }
        handler.handleEvents(eventList);
        Thread.sleep(500);
        // Send the second batch of DataMovementEvents
        eventList.clear();
        for (int j = 0; j < numOfPartitions - firstPart; j++) {
            Event dme = createDataMovementEvent(mapperHost, srcIndex++, targetIndex++);
            eventList.add(dme);
        }
        handler.handleEvents(eventList);
    }
    int waitCount = 100;
    while (waitCount-- > 0 && !(shuffleManager.isFetcherExecutorShutdown() && numOfMappers * numOfPartitions == shuffleManager.getNumOfCompletedInputs())) {
        Thread.sleep(100);
    }
    assertTrue(shuffleManager.isFetcherExecutorShutdown());
    assertEquals(numOfMappers * numOfPartitions, shuffleManager.getNumOfCompletedInputs());
}

Example 10 with FetchedInputAllocator

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

the class TestShuffleInputEventHandlerImpl method testCurrentPartitionEmpty.

@Test(timeout = 5000)
public void testCurrentPartitionEmpty() throws IOException {
    InputContext inputContext = mock(InputContext.class);
    ShuffleManager shuffleManager = mock(ShuffleManager.class);
    FetchedInputAllocator inputAllocator = mock(FetchedInputAllocator.class);
    ShuffleInputEventHandlerImpl handler = new ShuffleInputEventHandlerImpl(inputContext, shuffleManager, inputAllocator, null, false, 0, false);
    int taskIndex = 1;
    Event dme = createDataMovementEvent(0, taskIndex, createEmptyPartitionByteString(0));
    List<Event> eventList = new LinkedList<Event>();
    eventList.add(dme);
    handler.handleEvents(eventList);
    InputAttemptIdentifier expectedIdentifier = new InputAttemptIdentifier(taskIndex, 0);
    verify(shuffleManager).addCompletedInputWithNoData(eq(expectedIdentifier));
}