Search in sources :

Example 11 with DagInfo

use of org.apache.tez.serviceplugins.api.DagInfo in project tez by apache.

the class TestLocalTaskSchedulerService method preemptDescendantsOnly.

@Test
public void preemptDescendantsOnly() {
    final int MAX_TASKS = 2;
    TezConfiguration tezConf = new TezConfiguration();
    tezConf.setInt(TezConfiguration.TEZ_AM_INLINE_TASK_EXECUTION_MAX_TASKS, MAX_TASKS);
    ApplicationId appId = ApplicationId.newInstance(2000, 1);
    ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(appId, 1);
    Long parentTask1 = new Long(1);
    Long parentTask2 = new Long(2);
    Long childTask1 = new Long(3);
    Long grandchildTask1 = new Long(4);
    TaskSchedulerContext mockContext = TestTaskSchedulerHelpers.setupMockTaskSchedulerContext("", 0, "", true, appAttemptId, 1000l, null, tezConf);
    when(mockContext.getVertexIndexForTask(parentTask1)).thenReturn(0);
    when(mockContext.getVertexIndexForTask(parentTask2)).thenReturn(0);
    when(mockContext.getVertexIndexForTask(childTask1)).thenReturn(1);
    when(mockContext.getVertexIndexForTask(grandchildTask1)).thenReturn(2);
    DagInfo mockDagInfo = mock(DagInfo.class);
    when(mockDagInfo.getTotalVertices()).thenReturn(3);
    BitSet vertex1Descendants = new BitSet();
    vertex1Descendants.set(1);
    vertex1Descendants.set(2);
    BitSet vertex2Descendants = new BitSet();
    vertex2Descendants.set(2);
    BitSet vertex3Descendants = new BitSet();
    when(mockDagInfo.getVertexDescendants(0)).thenReturn(vertex1Descendants);
    when(mockDagInfo.getVertexDescendants(1)).thenReturn(vertex2Descendants);
    when(mockDagInfo.getVertexDescendants(2)).thenReturn(vertex3Descendants);
    when(mockContext.getCurrentDagInfo()).thenReturn(mockDagInfo);
    Priority priority1 = Priority.newInstance(1);
    Priority priority2 = Priority.newInstance(2);
    Priority priority3 = Priority.newInstance(3);
    Priority priority4 = Priority.newInstance(4);
    Resource resource = Resource.newInstance(1024, 1);
    MockLocalTaskSchedulerSerivce taskSchedulerService = new MockLocalTaskSchedulerSerivce(mockContext);
    // The mock context need to send a deallocate container request to the scheduler service
    Answer<Void> answer = new Answer<Void>() {

        @Override
        public Void answer(InvocationOnMock invocation) {
            ContainerId containerId = invocation.getArgumentAt(0, ContainerId.class);
            taskSchedulerService.deallocateContainer(containerId);
            return null;
        }
    };
    doAnswer(answer).when(mockContext).preemptContainer(any(ContainerId.class));
    taskSchedulerService.initialize();
    taskSchedulerService.start();
    taskSchedulerService.startRequestHandlerThread();
    MockAsyncDelegateRequestHandler requestHandler = taskSchedulerService.getRequestHandler();
    taskSchedulerService.allocateTask(parentTask1, resource, null, null, priority1, null, null);
    taskSchedulerService.allocateTask(childTask1, resource, null, null, priority3, null, null);
    taskSchedulerService.allocateTask(grandchildTask1, resource, null, null, priority4, null, null);
    requestHandler.drainRequest(3);
    // We should not preempt if we have not reached max task allocations
    Assert.assertEquals("Wrong number of allocate tasks", MAX_TASKS, requestHandler.allocateCount);
    Assert.assertTrue("Another allocation should not fit", !requestHandler.shouldProcess());
    // Next task allocation should preempt
    taskSchedulerService.allocateTask(parentTask2, Resource.newInstance(1024, 1), null, null, priority2, null, null);
    requestHandler.drainRequest(5);
    // All allocated tasks should have been removed
    Assert.assertEquals("Wrong number of preempted tasks", 1, requestHandler.preemptCount);
}
Also used : TaskSchedulerContext(org.apache.tez.serviceplugins.api.TaskSchedulerContext) DagInfo(org.apache.tez.serviceplugins.api.DagInfo) Priority(org.apache.hadoop.yarn.api.records.Priority) BitSet(java.util.BitSet) Resource(org.apache.hadoop.yarn.api.records.Resource) ApplicationAttemptId(org.apache.hadoop.yarn.api.records.ApplicationAttemptId) Answer(org.mockito.stubbing.Answer) ContainerId(org.apache.hadoop.yarn.api.records.ContainerId) InvocationOnMock(org.mockito.invocation.InvocationOnMock) ApplicationId(org.apache.hadoop.yarn.api.records.ApplicationId) MockAsyncDelegateRequestHandler(org.apache.tez.dag.app.rm.TestLocalTaskSchedulerService.MockLocalTaskSchedulerSerivce.MockAsyncDelegateRequestHandler) TezConfiguration(org.apache.tez.dag.api.TezConfiguration) Test(org.junit.Test)

Aggregations

BitSet (java.util.BitSet)11 DagInfo (org.apache.tez.serviceplugins.api.DagInfo)11 ApplicationAttemptId (org.apache.hadoop.yarn.api.records.ApplicationAttemptId)10 ContainerId (org.apache.hadoop.yarn.api.records.ContainerId)10 Priority (org.apache.hadoop.yarn.api.records.Priority)10 TezConfiguration (org.apache.tez.dag.api.TezConfiguration)10 TaskSchedulerContext (org.apache.tez.serviceplugins.api.TaskSchedulerContext)10 Test (org.junit.Test)10 Configuration (org.apache.hadoop.conf.Configuration)9 RegisterApplicationMasterResponse (org.apache.hadoop.yarn.api.protocolrecords.RegisterApplicationMasterResponse)9 Container (org.apache.hadoop.yarn.api.records.Container)9 NodeId (org.apache.hadoop.yarn.api.records.NodeId)9 MockClock (org.apache.tez.dag.app.MockClock)9 HeldContainer (org.apache.tez.dag.app.rm.DagAwareYarnTaskScheduler.HeldContainer)9 TaskRequest (org.apache.tez.dag.app.rm.DagAwareYarnTaskScheduler.TaskRequest)9 TaskSchedulerContextDrainable (org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.TaskSchedulerContextDrainable)9 TestTaskSchedulerHelpers.setupMockTaskSchedulerContext (org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.setupMockTaskSchedulerContext)9 AppFinalStatus (org.apache.tez.serviceplugins.api.TaskSchedulerContext.AppFinalStatus)8 ArrayList (java.util.ArrayList)3 Resource (org.apache.hadoop.yarn.api.records.Resource)2