Examples with TezCounters org.apache.tez.common.counters.TezCounters used on opensource projects

Example 26 with TezCounters

use of org.apache.tez.common.counters.TezCounters in project tez by apache.

the class TestShuffleUtils method createTezInputContext.

private InputContext createTezInputContext() {
    ApplicationId applicationId = ApplicationId.newInstance(1, 1);
    InputContext inputContext = mock(InputContext.class);
    doReturn(applicationId).when(inputContext).getApplicationId();
    doReturn("sourceVertex").when(inputContext).getSourceVertexName();
    when(inputContext.getCounters()).thenReturn(new TezCounters());
    return inputContext;
}

Example 27 with TezCounters

use of org.apache.tez.common.counters.TezCounters in project tez by apache.

the class TestShuffleUtils method createTezOutputContext.

private OutputContext createTezOutputContext() throws IOException {
    ApplicationId applicationId = ApplicationId.newInstance(1, 1);
    OutputContext outputContext = mock(OutputContext.class);
    ExecutionContextImpl executionContext = mock(ExecutionContextImpl.class);
    doReturn("localhost").when(executionContext).getHostName();
    doReturn(executionContext).when(outputContext).getExecutionContext();
    DataOutputBuffer serviceProviderMetaData = new DataOutputBuffer();
    serviceProviderMetaData.writeInt(80);
    doReturn(ByteBuffer.wrap(serviceProviderMetaData.getData())).when(outputContext).getServiceProviderMetaData(conf.get(TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID, TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID_DEFAULT));
    doReturn(1).when(outputContext).getTaskVertexIndex();
    doReturn(1).when(outputContext).getOutputIndex();
    doReturn(0).when(outputContext).getDAGAttemptNumber();
    doReturn("destVertex").when(outputContext).getDestinationVertexName();
    when(outputContext.getCounters()).thenReturn(new TezCounters());
    return outputContext;
}

Example 28 with TezCounters

use of org.apache.tez.common.counters.TezCounters in project tez by apache.

the class TestOrderedGroupedKVInput method testMergeConfig.

@Test
public void testMergeConfig() throws IOException, TezException {
    Configuration baseConf = new Configuration(false);
    baseConf.set("base-key", "base-value");
    Configuration payloadConf = new Configuration(false);
    payloadConf.set("local-key", "local-value");
    InputContext inputContext = mock(InputContext.class);
    UserPayload payLoad = TezUtils.createUserPayloadFromConf(payloadConf);
    String[] workingDirs = new String[] { "workDir1" };
    TezCounters counters = new TezCounters();
    doReturn(payLoad).when(inputContext).getUserPayload();
    doReturn(workingDirs).when(inputContext).getWorkDirs();
    doReturn(counters).when(inputContext).getCounters();
    doReturn(baseConf).when(inputContext).getContainerConfiguration();
    OrderedGroupedKVInput input = new OrderedGroupedKVInput(inputContext, 1);
    input.initialize();
    Configuration mergedConf = input.conf;
    assertEquals("base-value", mergedConf.get("base-key"));
    assertEquals("local-value", mergedConf.get("local-key"));
}

Example 29 with TezCounters

use of org.apache.tez.common.counters.TezCounters 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();
    doReturn(new Configuration(false)).when(inputContext).getContainerConfiguration();
    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 30 with TezCounters

use of org.apache.tez.common.counters.TezCounters in project tez by apache.

the class TestShuffle method createTezInputContext.

private InputContext createTezInputContext() throws IOException {
    ApplicationId applicationId = ApplicationId.newInstance(1, 1);
    InputContext inputContext = mock(InputContext.class);
    doReturn(applicationId).when(inputContext).getApplicationId();
    doReturn("sourceVertex").when(inputContext).getSourceVertexName();
    doReturn("taskVertex").when(inputContext).getTaskVertexName();
    when(inputContext.getCounters()).thenReturn(new TezCounters());
    ExecutionContext executionContext = new ExecutionContextImpl("localhost");
    doReturn(executionContext).when(inputContext).getExecutionContext();
    ByteBuffer shuffleBuffer = ByteBuffer.allocate(4).putInt(0, 4);
    doReturn(shuffleBuffer).when(inputContext).getServiceProviderMetaData(anyString());
    Token<JobTokenIdentifier> sessionToken = new Token<JobTokenIdentifier>(new JobTokenIdentifier(new Text("text")), new JobTokenSecretManager());
    ByteBuffer tokenBuffer = TezCommonUtils.serializeServiceData(sessionToken);
    doReturn(tokenBuffer).when(inputContext).getServiceConsumerMetaData(anyString());
    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;
}