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

Example 16 with InputContext

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

the class TestShuffleInputEventHandlerImpl method testPipelinedShuffleEvents_WithEmptyPartitions.

/**
 * In pipelined shuffle, check if processing & exceptions are done correctly when empty
 * partitions are sent
 *
 * @throws IOException
 */
@Test(timeout = 5000)
public void testPipelinedShuffleEvents_WithEmptyPartitions() throws IOException {
    InputContext inputContext = createInputContext();
    ShuffleManager shuffleManager = createShuffleManager(inputContext);
    FetchedInputAllocator inputAllocator = mock(FetchedInputAllocator.class);
    ShuffleInputEventHandlerImpl handler = new ShuffleInputEventHandlerImpl(inputContext, shuffleManager, inputAllocator, null, false, 0, false);
    // 0--> 1 with spill id 0 (attemptNum 0) with empty partitions
    BitSet bitSet = new BitSet(4);
    bitSet.flip(0, 4);
    Event dme = createDataMovementEvent(true, 0, 1, 0, false, bitSet, 4, 0);
    handler.handleEvents(Collections.singletonList(dme));
    InputAttemptIdentifier expected = new InputAttemptIdentifier(1, 0, PATH_COMPONENT, false, InputAttemptIdentifier.SPILL_INFO.INCREMENTAL_UPDATE, 0);
    verify(shuffleManager, times(1)).addCompletedInputWithNoData(expected);
    // 0--> 1 with spill id 1 (attemptNum 0)
    handler.handleEvents(Collections.singletonList(dme));
    dme = createDataMovementEvent(true, 0, 1, 1, false, new BitSet(), 4, 0);
    expected = new InputAttemptIdentifier(1, 0, PATH_COMPONENT, false, InputAttemptIdentifier.SPILL_INFO.INCREMENTAL_UPDATE, 1);
    verify(shuffleManager, times(2)).addCompletedInputWithNoData(expected);
    // Let attemptNum 0 be scheduled.
    shuffleManager.shuffleInfoEventsMap.get(expected.getInputIdentifier()).scheduledForDownload = true;
    // Now send attemptNum 1.  This should throw exception, because attempt #1 is already added
    dme = createDataMovementEvent(true, 0, 1, 0, false, new BitSet(), 4, 1);
    handler.handleEvents(Collections.singletonList(dme));
    verify(inputContext).killSelf(any(Throwable.class), anyString());
}

Example 17 with InputContext

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

the class TestShuffleInputEventHandlerImpl method testSimple.

@Test(timeout = 5000)
public void testSimple() 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, null);
    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 18 with InputContext

use of org.apache.tez.runtime.api.InputContext 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 19 with InputContext

use of org.apache.tez.runtime.api.InputContext 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 20 with InputContext

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

the class TestShuffleInputEventHandlerImpl method createInputContext.

private InputContext createInputContext() throws IOException {
    DataOutputBuffer port_dob = new DataOutputBuffer();
    port_dob.writeInt(PORT);
    final ByteBuffer shuffleMetaData = ByteBuffer.wrap(port_dob.getData(), 0, port_dob.getLength());
    port_dob.close();
    ExecutionContext executionContext = mock(ExecutionContext.class);
    doReturn(HOST).when(executionContext).getHostName();
    InputContext inputContext = mock(InputContext.class);
    doReturn(new TezCounters()).when(inputContext).getCounters();
    doReturn("sourceVertex").when(inputContext).getSourceVertexName();
    doReturn(shuffleMetaData).when(inputContext).getServiceProviderMetaData(conf.get(TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID, TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID_DEFAULT));
    doReturn(executionContext).when(inputContext).getExecutionContext();
    when(inputContext.createTezFrameworkExecutorService(anyInt(), anyString())).thenAnswer(new Answer<ExecutorService>() {

        @Override
        public ExecutorService answer(InvocationOnMock invocation) throws Throwable {
            return sharedExecutor.createExecutorService(invocation.getArgumentAt(0, Integer.class), invocation.getArgumentAt(1, String.class));
        }
    });
    return inputContext;
}