use of org.apache.tez.runtime.api.impl.TezUmbilical 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;
}
use of org.apache.tez.runtime.api.impl.TezUmbilical in project tez by apache.
the class TestLogicalIOProcessorRuntimeTask method testAutoStart.
@Test(timeout = 5000)
public void testAutoStart() throws Exception {
TezDAGID dagId = createTezDagId();
TezVertexID vertexId = createTezVertexId(dagId);
Map<String, ByteBuffer> serviceConsumerMetadata = new HashMap<String, ByteBuffer>();
Multimap<String, String> startedInputsMap = HashMultimap.create();
TezUmbilical umbilical = mock(TezUmbilical.class);
TezConfiguration tezConf = new TezConfiguration();
tezConf.set(TezConfiguration.TEZ_TASK_SCALE_MEMORY_ALLOCATOR_CLASS, ScalingAllocator.class.getName());
TezTaskAttemptID taId1 = createTaskAttemptID(vertexId, 1);
TaskSpec task1 = createTaskSpec(taId1, "dag1", "vertex1", 30);
TezTaskAttemptID taId2 = createTaskAttemptID(vertexId, 2);
TaskSpec task2 = createTaskSpec(taId2, "dag2", "vertex1", 10);
TezSharedExecutor sharedExecutor = new TezSharedExecutor(tezConf);
LogicalIOProcessorRuntimeTask lio1 = new LogicalIOProcessorRuntimeTask(task1, 0, tezConf, null, umbilical, serviceConsumerMetadata, new HashMap<String, String>(), startedInputsMap, null, "", new ExecutionContextImpl("localhost"), Runtime.getRuntime().maxMemory(), true, new DefaultHadoopShim(), sharedExecutor);
try {
lio1.initialize();
lio1.run();
lio1.close();
// Input should've been started, Output should not have been started
assertEquals(1, TestProcessor.runCount);
assertEquals(1, TestInput.startCount);
assertEquals(0, TestOutput.startCount);
// test that invocations of progress are counted correctly
assertEquals(true, lio1.getAndClearProgressNotification());
// cleared after getting
assertEquals(false, lio1.getAndClearProgressNotification());
assertEquals(30, TestInput.vertexParallelism);
assertEquals(0, TestOutput.vertexParallelism);
assertEquals(30, lio1.getProcessorContext().getVertexParallelism());
assertEquals(30, lio1.getInputContexts().iterator().next().getVertexParallelism());
assertEquals(30, lio1.getOutputContexts().iterator().next().getVertexParallelism());
} catch (Exception e) {
fail();
sharedExecutor.shutdownNow();
} finally {
cleanupAndTest(lio1);
}
// local mode
tezConf.setBoolean(TezConfiguration.TEZ_LOCAL_MODE, true);
LogicalIOProcessorRuntimeTask lio2 = new LogicalIOProcessorRuntimeTask(task2, 0, tezConf, null, umbilical, serviceConsumerMetadata, new HashMap<String, String>(), startedInputsMap, null, "", new ExecutionContextImpl("localhost"), Runtime.getRuntime().maxMemory(), true, new DefaultHadoopShim(), sharedExecutor);
try {
lio2.initialize();
lio2.run();
lio2.close();
// Input should not have been started again, Output should not have been started
assertEquals(2, TestProcessor.runCount);
assertEquals(1, TestInput.startCount);
assertEquals(0, TestOutput.startCount);
assertEquals(30, TestInput.vertexParallelism);
assertEquals(0, TestOutput.vertexParallelism);
// Check if parallelism is available in processor/ i/p / o/p contexts
assertEquals(10, lio2.getProcessorContext().getVertexParallelism());
assertEquals(10, lio2.getInputContexts().iterator().next().getVertexParallelism());
assertEquals(10, lio2.getOutputContexts().iterator().next().getVertexParallelism());
} catch (Exception e) {
fail();
} finally {
cleanupAndTest(lio2);
sharedExecutor.shutdownNow();
}
}
Aggregations