use of org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.TaskSchedulerWithDrainableContext in project tez by apache.
the class TestContainerReuse method testDifferentResourceContainerReuse.
@Test(timeout = 5000)
public void testDifferentResourceContainerReuse() throws Exception {
Configuration tezConf = new Configuration(new YarnConfiguration());
tezConf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
tezConf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, true);
tezConf.setLong(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 0);
tezConf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MIN_MILLIS, 0);
tezConf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MAX_MILLIS, 0);
CapturingEventHandler eventHandler = new CapturingEventHandler();
TezDAGID dagID = TezDAGID.getInstance("0", 0, 0);
AMRMClient<CookieContainerRequest> rmClientCore = new AMRMClientForTest();
TezAMRMClientAsync<CookieContainerRequest> rmClient = spy(new AMRMClientAsyncForTest(rmClientCore, 100));
AppContext appContext = mock(AppContext.class);
doReturn(new Configuration(false)).when(appContext).getAMConf();
AMContainerMap amContainerMap = new AMContainerMap(mock(ContainerHeartbeatHandler.class), mock(TaskCommunicatorManagerInterface.class), new ContainerContextMatcher(), appContext);
AMNodeTracker amNodeTracker = new AMNodeTracker(eventHandler, appContext);
doReturn(amContainerMap).when(appContext).getAllContainers();
doReturn(amNodeTracker).when(appContext).getNodeTracker();
doReturn(DAGAppMasterState.RUNNING).when(appContext).getAMState();
doReturn(dagID).when(appContext).getCurrentDAGID();
doReturn(mock(ClusterInfo.class)).when(appContext).getClusterInfo();
TaskSchedulerManager taskSchedulerManagerReal = new TaskSchedulerManagerForTest(appContext, eventHandler, rmClient, new AlwaysMatchesContainerMatcher(), TezUtils.createUserPayloadFromConf(tezConf));
TaskSchedulerManager taskSchedulerManager = spy(taskSchedulerManagerReal);
taskSchedulerManager.init(tezConf);
taskSchedulerManager.start();
TaskSchedulerWithDrainableContext taskScheduler = (TaskSchedulerWithDrainableContext) ((TaskSchedulerManagerForTest) taskSchedulerManager).getSpyTaskScheduler();
TaskSchedulerContextDrainable drainableAppCallback = taskScheduler.getDrainableAppCallback();
AtomicBoolean drainNotifier = new AtomicBoolean(false);
taskScheduler.delayedContainerManager.drainedDelayedContainersForTest = drainNotifier;
Resource resource1 = Resource.newInstance(1024, 1);
Resource resource2 = Resource.newInstance(2048, 1);
String[] host1 = { "host1" };
String[] host2 = { "host2" };
String[] racks = { "/default-rack" };
Priority priority1 = Priority.newInstance(1);
TezVertexID vertexID1 = TezVertexID.getInstance(dagID, 1);
// Vertex 1, Task 1, Attempt 1, host1
TezTaskAttemptID taID11 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID1, 1), 1);
TaskAttempt ta11 = mock(TaskAttempt.class);
AMSchedulerEventTALaunchRequest lrEvent1 = createLaunchRequestEvent(taID11, ta11, resource1, host1, racks, priority1);
// Vertex 1, Task 2, Attempt 1, host1
TezTaskAttemptID taID12 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID1, 2), 1);
TaskAttempt ta12 = mock(TaskAttempt.class);
AMSchedulerEventTALaunchRequest lrEvent2 = createLaunchRequestEvent(taID12, ta12, resource1, host1, racks, priority1);
// Vertex 1, Task 3, Attempt 1, host2
TezTaskAttemptID taID13 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID1, 3), 1);
TaskAttempt ta13 = mock(TaskAttempt.class);
AMSchedulerEventTALaunchRequest lrEvent3 = createLaunchRequestEvent(taID13, ta13, resource2, host2, racks, priority1);
// Vertex 1, Task 4, Attempt 1, host2
TezTaskAttemptID taID14 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID1, 4), 1);
TaskAttempt ta14 = mock(TaskAttempt.class);
AMSchedulerEventTALaunchRequest lrEvent4 = createLaunchRequestEvent(taID14, ta14, resource2, host2, racks, priority1);
taskSchedulerManager.handleEvent(lrEvent1);
taskSchedulerManager.handleEvent(lrEvent2);
taskSchedulerManager.handleEvent(lrEvent3);
taskSchedulerManager.handleEvent(lrEvent4);
Container container1 = createContainer(1, "host1", resource1, priority1);
Container container2 = createContainer(2, "host2", resource2, priority1);
// One container allocated, should start ta11
taskScheduler.onContainersAllocated(Collections.singletonList(container1));
TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier);
drainableAppCallback.drain();
verify(taskSchedulerManager).taskAllocated(eq(0), eq(ta11), any(Object.class), eq(container1));
// Second container allocated, should start ta13
taskScheduler.onContainersAllocated(Collections.singletonList(container2));
TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier);
drainableAppCallback.drain();
verify(taskSchedulerManager).taskAllocated(eq(0), eq(ta13), any(Object.class), eq(container2));
// ta11 finished, should start ta12
taskSchedulerManager.handleEvent(new AMSchedulerEventTAEnded(ta11, container1.getId(), TaskAttemptState.SUCCEEDED, null, null, 0));
drainableAppCallback.drain();
verifyDeAllocateTask(taskScheduler, ta11, true, null, null);
verify(taskSchedulerManager).taskAllocated(eq(0), eq(ta12), any(Object.class), eq(container1));
verify(rmClient, times(0)).releaseAssignedContainer(eq(container1.getId()));
eventHandler.verifyNoInvocations(AMContainerEventStopRequest.class);
eventHandler.reset();
// ta13 finished, should start ta14
taskSchedulerManager.handleEvent(new AMSchedulerEventTAEnded(ta13, container2.getId(), TaskAttemptState.SUCCEEDED, null, null, 0));
drainableAppCallback.drain();
verifyDeAllocateTask(taskScheduler, ta13, true, null, null);
verify(taskSchedulerManager).taskAllocated(eq(0), eq(ta14), any(Object.class), eq(container2));
verify(rmClient, times(0)).releaseAssignedContainer(eq(container2.getId()));
eventHandler.verifyNoInvocations(AMContainerEventStopRequest.class);
eventHandler.reset();
// ta12 finished no pending requests, should release container1
taskSchedulerManager.handleEvent(new AMSchedulerEventTAEnded(ta12, container1.getId(), TaskAttemptState.SUCCEEDED, null, null, 0));
drainableAppCallback.drain();
verifyDeAllocateTask(taskScheduler, ta12, true, null, null);
verify(rmClient).releaseAssignedContainer(eq(container1.getId()));
eventHandler.verifyInvocation(AMContainerEventStopRequest.class);
eventHandler.reset();
// ta14 finished no pending requests, should release container2.
taskSchedulerManager.handleEvent(new AMSchedulerEventTAEnded(ta14, container2.getId(), TaskAttemptState.SUCCEEDED, null, null, 0));
drainableAppCallback.drain();
verifyDeAllocateTask(taskScheduler, ta14, true, null, null);
verify(rmClient).releaseAssignedContainer(eq(container2.getId()));
eventHandler.verifyInvocation(AMContainerEventStopRequest.class);
eventHandler.reset();
taskScheduler.shutdown();
taskSchedulerManager.close();
}
use of org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.TaskSchedulerWithDrainableContext in project tez by apache.
the class TestTaskScheduler method testLocalityMatching.
@Test(timeout = 5000)
public void testLocalityMatching() throws Exception {
TezAMRMClientAsync<CookieContainerRequest> amrmClient = spy(new AMRMClientAsyncForTest(new AMRMClientForTest(), 100));
Configuration conf = new Configuration();
conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, false);
TaskSchedulerContext appClient = setupMockTaskSchedulerContext(DEFAULT_APP_HOST, DEFAULT_APP_PORT, "", conf);
final TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(appClient);
TaskSchedulerWithDrainableContext taskScheduler = new TaskSchedulerWithDrainableContext(drainableAppCallback, amrmClient);
taskScheduler.initialize();
taskScheduler.start();
Resource resource = Resource.newInstance(1024, 1);
Priority priority = Priority.newInstance(1);
String[] hostsTask1 = { "host1" };
String[] hostsTask2 = { "non-allocated-host" };
String[] defaultRack = { "/default-rack" };
String[] otherRack = { "/other-rack" };
Object mockTask1 = mock(Object.class);
CookieContainerRequest mockCookie1 = mock(CookieContainerRequest.class, RETURNS_DEEP_STUBS);
when(mockCookie1.getCookie().getTask()).thenReturn(mockTask1);
Object mockTask2 = mock(Object.class);
CookieContainerRequest mockCookie2 = mock(CookieContainerRequest.class, RETURNS_DEEP_STUBS);
when(mockCookie2.getCookie().getTask()).thenReturn(mockTask2);
Container containerHost1 = createContainer(1, "host1", resource, priority);
Container containerHost3 = createContainer(2, "host3", resource, priority);
List<Container> allocatedContainers = new LinkedList<Container>();
allocatedContainers.add(containerHost3);
allocatedContainers.add(containerHost1);
taskScheduler.allocateTask(mockTask1, resource, hostsTask1, defaultRack, priority, null, mockCookie1);
drainableAppCallback.drain();
List<CookieContainerRequest> host1List = new ArrayList<CookieContainerRequest>();
host1List.add(mockCookie1);
List<CookieContainerRequest> defaultRackList = new ArrayList<CookieContainerRequest>();
defaultRackList.add(mockCookie1);
List<CookieContainerRequest> nonAllocatedHostList = new ArrayList<YarnTaskSchedulerService.CookieContainerRequest>();
nonAllocatedHostList.add(mockCookie2);
List<CookieContainerRequest> otherRackList = new ArrayList<YarnTaskSchedulerService.CookieContainerRequest>();
otherRackList.add(mockCookie2);
taskScheduler.allocateTask(mockTask2, resource, hostsTask2, otherRack, priority, null, mockCookie2);
drainableAppCallback.drain();
List<CookieContainerRequest> anyList = new LinkedList<YarnTaskSchedulerService.CookieContainerRequest>();
anyList.add(mockCookie1);
anyList.add(mockCookie2);
taskScheduler.onContainersAllocated(allocatedContainers);
drainableAppCallback.drain();
ArgumentCaptor<Object> taskCaptor = ArgumentCaptor.forClass(Object.class);
ArgumentCaptor<Container> containerCaptor = ArgumentCaptor.forClass(Container.class);
verify(appClient, times(2)).taskAllocated(taskCaptor.capture(), any(), containerCaptor.capture());
// Expected containerHost1 allocated to task1 due to locality,
// containerHost3 allocated to task2.
List<Container> assignedContainers = containerCaptor.getAllValues();
int container1Pos = assignedContainers.indexOf(containerHost1);
assertTrue("Container: " + containerHost1 + " was not assigned", container1Pos != -1);
assertEquals("Task 1 was not allocated to containerHost1", mockTask1, taskCaptor.getAllValues().get(container1Pos));
int container2Pos = assignedContainers.indexOf(containerHost3);
assertTrue("Container: " + containerHost3 + " was not assigned", container2Pos != -1);
assertEquals("Task 2 was not allocated to containerHost3", mockTask2, taskCaptor.getAllValues().get(container2Pos));
AppFinalStatus finalStatus = new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, "", "");
when(appClient.getFinalAppStatus()).thenReturn(finalStatus);
taskScheduler.shutdown();
}
use of org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.TaskSchedulerWithDrainableContext in project tez by apache.
the class TestTaskScheduler method testTaskSchedulerPreemption2.
@Test(timeout = 5000)
public void testTaskSchedulerPreemption2() throws Exception {
TezAMRMClientAsync<CookieContainerRequest> mockRMClient = spy(new AMRMClientAsyncForTest(new AMRMClientForTest(), 100));
int waitTime = 1000;
Configuration conf = new Configuration();
conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, false);
conf.setInt(TezConfiguration.TEZ_AM_PREEMPTION_HEARTBEATS_BETWEEN_PREEMPTIONS, 2);
conf.setInt(TezConfiguration.TEZ_AM_PREEMPTION_MAX_WAIT_TIME_MS, waitTime);
TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(DEFAULT_APP_HOST, DEFAULT_APP_PORT, DEFAULT_APP_URL, false, null, null, new PreemptionMatcher(), conf);
final TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);
final TaskSchedulerWithDrainableContext scheduler = new TaskSchedulerWithDrainableContext(drainableAppCallback, mockRMClient);
scheduler.initialize();
scheduler.start();
// no preemption
scheduler.getProgress();
drainableAppCallback.drain();
Resource totalResource = mockRMClient.getAvailableResources();
Assert.assertEquals(totalResource, scheduler.getTotalResources());
verify(mockRMClient, times(0)).releaseAssignedContainer((ContainerId) any());
// allocate task
Object mockTask1 = mock(Object.class);
Object mockTask2 = mock(Object.class);
Object mockTask3 = mock(Object.class);
Object obj3 = new Object();
Priority pri2 = Priority.newInstance(2);
Priority pri4 = Priority.newInstance(4);
Priority pri6 = Priority.newInstance(6);
ArgumentCaptor<CookieContainerRequest> requestCaptor = ArgumentCaptor.forClass(CookieContainerRequest.class);
final ArrayList<CookieContainerRequest> anyContainers = new ArrayList<CookieContainerRequest>();
Resource taskAsk = Resource.newInstance(1024, 1);
scheduler.allocateTask(mockTask1, taskAsk, null, null, pri4, null, null);
drainableAppCallback.drain();
verify(mockRMClient, times(1)).addContainerRequest(requestCaptor.capture());
anyContainers.add(requestCaptor.getValue());
scheduler.getProgress();
drainableAppCallback.drain();
Assert.assertEquals(totalResource, scheduler.getTotalResources());
verify(mockRMClient, times(0)).releaseAssignedContainer((ContainerId) any());
List<Container> containers = new ArrayList<Container>();
Container mockContainer1 = mock(Container.class, RETURNS_DEEP_STUBS);
when(mockContainer1.getNodeId().getHost()).thenReturn("host1");
when(mockContainer1.getResource()).thenReturn(taskAsk);
when(mockContainer1.getPriority()).thenReturn(pri4);
ContainerId mockCId1 = mock(ContainerId.class);
when(mockContainer1.getId()).thenReturn(mockCId1);
containers.add(mockContainer1);
Mockito.doAnswer(new Answer<Object>() {
public Object answer(InvocationOnMock invocation) {
Object[] args = invocation.getArguments();
ContainerId cId = (ContainerId) args[0];
scheduler.deallocateContainer(cId);
return null;
}
}).when(mockApp).preemptContainer((ContainerId) any());
scheduler.onContainersAllocated(containers);
drainableAppCallback.drain();
Assert.assertEquals(1, scheduler.taskAllocations.size());
Assert.assertEquals(mockCId1, scheduler.taskAllocations.get(mockTask1).getId());
// no preemption
scheduler.getProgress();
drainableAppCallback.drain();
verify(mockRMClient, times(0)).releaseAssignedContainer((ContainerId) any());
// no need for task preemption until now - so they should match
Assert.assertEquals(scheduler.numHeartbeats, scheduler.heartbeatAtLastPreemption);
// add a pending request that cannot be allocated until resources free up
Object mockTask2Cookie = new Object();
scheduler.allocateTask(mockTask2, taskAsk, null, null, pri2, obj3, mockTask2Cookie);
Object mockTask3Cookie = new Object();
scheduler.allocateTask(mockTask3, taskAsk, null, null, pri6, obj3, mockTask3Cookie);
// nothing waiting till now
Assert.assertNull(scheduler.highestWaitingRequestPriority);
Assert.assertEquals(0, scheduler.highestWaitingRequestWaitStartTime);
long currTime = System.currentTimeMillis();
scheduler.getProgress();
drainableAppCallback.drain();
verify(mockRMClient, times(0)).releaseAssignedContainer((ContainerId) any());
// enough free resources. preemption not triggered
Assert.assertEquals(pri2, scheduler.highestWaitingRequestPriority);
Assert.assertTrue(scheduler.highestWaitingRequestWaitStartTime >= currTime);
Assert.assertEquals(scheduler.numHeartbeats, scheduler.heartbeatAtLastPreemption);
Thread.sleep(waitTime + 10);
long oldStartWaitTime = scheduler.highestWaitingRequestWaitStartTime;
scheduler.getProgress();
drainableAppCallback.drain();
verify(mockRMClient, times(0)).releaseAssignedContainer((ContainerId) any());
// enough free resources. deadline crossed. preemption triggered
Assert.assertEquals(pri2, scheduler.highestWaitingRequestPriority);
Assert.assertEquals(oldStartWaitTime, scheduler.highestWaitingRequestWaitStartTime);
Assert.assertTrue(scheduler.numHeartbeats > scheduler.heartbeatAtLastPreemption);
scheduler.getProgress();
drainableAppCallback.drain();
verify(mockRMClient, times(1)).releaseAssignedContainer((ContainerId) any());
verify(mockRMClient, times(1)).releaseAssignedContainer(mockCId1);
Assert.assertEquals(scheduler.numHeartbeats, scheduler.heartbeatAtLastPreemption);
// maintains existing waiting values
Assert.assertEquals(pri2, scheduler.highestWaitingRequestPriority);
Assert.assertEquals(oldStartWaitTime, scheduler.highestWaitingRequestWaitStartTime);
// remove high pri request to test waiting pri change
scheduler.deallocateTask(mockTask2, false, null, null);
scheduler.getProgress();
// waiting value changes
Assert.assertEquals(pri6, scheduler.highestWaitingRequestPriority);
Assert.assertTrue(oldStartWaitTime < scheduler.highestWaitingRequestWaitStartTime);
Thread.sleep(waitTime + 10);
scheduler.getProgress();
drainableAppCallback.drain();
// deadlines crossed but nothing lower pri running. so reset
Assert.assertNull(scheduler.highestWaitingRequestPriority);
Assert.assertEquals(0, scheduler.highestWaitingRequestWaitStartTime);
scheduler.getProgress();
drainableAppCallback.drain();
// waiting value changes
Assert.assertEquals(pri6, scheduler.highestWaitingRequestPriority);
Assert.assertTrue(oldStartWaitTime < scheduler.highestWaitingRequestWaitStartTime);
AppFinalStatus finalStatus = new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, "", DEFAULT_APP_URL);
when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
scheduler.shutdown();
drainableAppCallback.drain();
}
use of org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.TaskSchedulerWithDrainableContext in project tez by apache.
the class TestTaskScheduler method testTaskSchedulerPreemption.
@SuppressWarnings({ "unchecked", "rawtypes" })
@Test(timeout = 5000)
public void testTaskSchedulerPreemption() throws Exception {
TezAMRMClientAsync<CookieContainerRequest> mockRMClient = mock(TezAMRMClientAsync.class);
Configuration conf = new Configuration();
conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, false);
conf.setInt(TezConfiguration.TEZ_AM_PREEMPTION_HEARTBEATS_BETWEEN_PREEMPTIONS, 3);
TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(DEFAULT_APP_HOST, DEFAULT_APP_PORT, DEFAULT_APP_URL, false, null, null, new PreemptionMatcher(), conf);
final TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);
final TaskSchedulerWithDrainableContext scheduler = new TaskSchedulerWithDrainableContext(drainableAppCallback, mockRMClient);
scheduler.initialize();
RegisterApplicationMasterResponse mockRegResponse = mock(RegisterApplicationMasterResponse.class);
when(mockRMClient.registerApplicationMaster(anyString(), anyInt(), anyString())).thenReturn(mockRegResponse);
scheduler.start();
Resource totalResource = Resource.newInstance(4000, 4);
when(mockRMClient.getAvailableResources()).thenReturn(totalResource);
// no preemption
scheduler.getProgress();
drainableAppCallback.drain();
Assert.assertEquals(totalResource, scheduler.getTotalResources());
verify(mockRMClient, times(0)).releaseAssignedContainer((ContainerId) any());
// allocate task
Object mockTask1 = mock(Object.class);
Object mockTask2 = mock(Object.class);
Object mockTask3 = mock(Object.class);
Object mockTask3Wait = mock(Object.class);
Object mockTask3Retry = mock(Object.class);
Object mockTask3KillA = mock(Object.class);
Object mockTask3KillB = mock(Object.class);
Object mockTaskPri8 = mock(Object.class);
Object obj3 = new Object();
Priority pri2 = Priority.newInstance(2);
Priority pri4 = Priority.newInstance(4);
Priority pri5 = Priority.newInstance(5);
Priority pri6 = Priority.newInstance(6);
Priority pri8 = Priority.newInstance(8);
ArgumentCaptor<CookieContainerRequest> requestCaptor = ArgumentCaptor.forClass(CookieContainerRequest.class);
final ArrayList<CookieContainerRequest> anyContainers = new ArrayList<CookieContainerRequest>();
Resource taskAsk = Resource.newInstance(1024, 1);
scheduler.allocateTask(mockTask1, taskAsk, null, null, pri2, null, null);
drainableAppCallback.drain();
verify(mockRMClient, times(1)).addContainerRequest(requestCaptor.capture());
anyContainers.add(requestCaptor.getValue());
scheduler.allocateTask(mockTask3, taskAsk, null, null, pri6, obj3, null);
drainableAppCallback.drain();
verify(mockRMClient, times(2)).addContainerRequest(requestCaptor.capture());
anyContainers.add(requestCaptor.getValue());
// later one in the allocation gets killed between the two task3's
scheduler.allocateTask(mockTask3KillA, taskAsk, null, null, pri6, obj3, null);
drainableAppCallback.drain();
verify(mockRMClient, times(3)).addContainerRequest(requestCaptor.capture());
anyContainers.add(requestCaptor.getValue());
// later one in the allocation gets killed between the two task3's
scheduler.allocateTask(mockTask3KillB, taskAsk, null, null, pri6, obj3, null);
drainableAppCallback.drain();
verify(mockRMClient, times(4)).addContainerRequest(requestCaptor.capture());
anyContainers.add(requestCaptor.getValue());
Resource freeResource = Resource.newInstance(500, 0);
when(mockRMClient.getAvailableResources()).thenReturn(freeResource);
scheduler.getProgress();
drainableAppCallback.drain();
Assert.assertEquals(totalResource, scheduler.getTotalResources());
verify(mockRMClient, times(0)).releaseAssignedContainer((ContainerId) any());
final List<ArrayList<CookieContainerRequest>> anyList = new LinkedList<ArrayList<CookieContainerRequest>>();
final List<ArrayList<CookieContainerRequest>> emptyList = new LinkedList<ArrayList<CookieContainerRequest>>();
anyList.add(anyContainers);
List<Container> containers = new ArrayList<Container>();
Container mockContainer1 = mock(Container.class, RETURNS_DEEP_STUBS);
when(mockContainer1.getNodeId().getHost()).thenReturn("host1");
when(mockContainer1.getResource()).thenReturn(taskAsk);
when(mockContainer1.getPriority()).thenReturn(pri2);
ContainerId mockCId1 = mock(ContainerId.class);
when(mockContainer1.getId()).thenReturn(mockCId1);
containers.add(mockContainer1);
Container mockContainer2 = mock(Container.class, RETURNS_DEEP_STUBS);
when(mockContainer2.getNodeId().getHost()).thenReturn("host1");
when(mockContainer2.getResource()).thenReturn(taskAsk);
when(mockContainer2.getPriority()).thenReturn(pri6);
ContainerId mockCId2 = mock(ContainerId.class);
when(mockContainer2.getId()).thenReturn(mockCId2);
containers.add(mockContainer2);
Container mockContainer3A = mock(Container.class, RETURNS_DEEP_STUBS);
when(mockContainer3A.getNodeId().getHost()).thenReturn("host1");
when(mockContainer3A.getResource()).thenReturn(taskAsk);
when(mockContainer3A.getPriority()).thenReturn(pri6);
ContainerId mockCId3A = mock(ContainerId.class);
when(mockContainer3A.getId()).thenReturn(mockCId3A);
containers.add(mockContainer3A);
Container mockContainer3B = mock(Container.class, RETURNS_DEEP_STUBS);
when(mockContainer3B.getNodeId().getHost()).thenReturn("host1");
when(mockContainer3B.getResource()).thenReturn(taskAsk);
// high priority container
when(mockContainer3B.getPriority()).thenReturn(pri2);
ContainerId mockCId3B = mock(ContainerId.class);
when(mockContainer3B.getId()).thenReturn(mockCId3B);
containers.add(mockContainer3B);
when(mockRMClient.getMatchingRequests((Priority) any(), eq("host1"), (Resource) any())).thenAnswer(new Answer<List<? extends Collection<CookieContainerRequest>>>() {
@Override
public List<? extends Collection<CookieContainerRequest>> answer(InvocationOnMock invocation) throws Throwable {
return emptyList;
}
});
// RackResolver by default puts hosts in default-rack
when(mockRMClient.getMatchingRequests((Priority) any(), eq("/default-rack"), (Resource) any())).thenAnswer(new Answer<List<? extends Collection<CookieContainerRequest>>>() {
@Override
public List<? extends Collection<CookieContainerRequest>> answer(InvocationOnMock invocation) throws Throwable {
return emptyList;
}
});
when(mockRMClient.getMatchingRequests((Priority) any(), eq(ResourceRequest.ANY), (Resource) any())).thenAnswer(new Answer<List<? extends Collection<CookieContainerRequest>>>() {
int calls = 0;
@Override
public List<? extends Collection<CookieContainerRequest>> answer(InvocationOnMock invocation) throws Throwable {
if (calls > 0) {
anyContainers.remove(0);
}
calls++;
return anyList;
}
});
Mockito.doAnswer(new Answer() {
public Object answer(InvocationOnMock invocation) {
Object[] args = invocation.getArguments();
ContainerId cId = (ContainerId) args[0];
scheduler.deallocateContainer(cId);
return null;
}
}).when(mockApp).preemptContainer((ContainerId) any());
scheduler.onContainersAllocated(containers);
drainableAppCallback.drain();
Assert.assertEquals(4, scheduler.taskAllocations.size());
Assert.assertEquals(4096, scheduler.allocatedResources.getMemory());
Assert.assertEquals(mockCId1, scheduler.taskAllocations.get(mockTask1).getId());
Assert.assertEquals(mockCId2, scheduler.taskAllocations.get(mockTask3).getId());
Assert.assertEquals(mockCId3A, scheduler.taskAllocations.get(mockTask3KillA).getId());
// high priority container assigned to lower pri task. This task should still be preempted
// because the task priority is relevant for preemption and not the container priority
Assert.assertEquals(mockCId3B, scheduler.taskAllocations.get(mockTask3KillB).getId());
// no preemption
scheduler.getProgress();
drainableAppCallback.drain();
verify(mockRMClient, times(0)).releaseAssignedContainer((ContainerId) any());
// no need for task preemption until now - so they should match
Assert.assertEquals(scheduler.numHeartbeats, scheduler.heartbeatAtLastPreemption);
// add a pending request that cannot be allocated until resources free up
Object mockTask3WaitCookie = new Object();
scheduler.allocateTask(mockTask3Wait, taskAsk, null, null, pri6, obj3, mockTask3WaitCookie);
// add a pri 8 request for the pri 8 container that will not be matched
Object mockTaskPri8Cookie = new Object();
scheduler.allocateTask(mockTaskPri8, taskAsk, null, null, pri8, obj3, mockTaskPri8Cookie);
// no preemption - same pri
scheduler.getProgress();
drainableAppCallback.drain();
verify(mockRMClient, times(6)).addContainerRequest(requestCaptor.capture());
verify(mockRMClient, times(0)).releaseAssignedContainer((ContainerId) any());
Container mockContainer4 = mock(Container.class, RETURNS_DEEP_STUBS);
when(mockContainer4.getNodeId().getHost()).thenReturn("host1");
when(mockContainer4.getResource()).thenReturn(taskAsk);
when(mockContainer4.getPriority()).thenReturn(pri8);
ContainerId mockCId4 = mock(ContainerId.class);
when(mockContainer4.getId()).thenReturn(mockCId4);
containers.clear();
containers.add(mockContainer4);
// new lower pri container added that wont be matched and eventually preempted
// Fudge new container being present in delayed allocation list due to race
HeldContainer heldContainer = new HeldContainer(mockContainer4, -1, -1, null, containerSignatureMatcher);
scheduler.delayedContainerManager.delayedContainers.add(heldContainer);
// no preemption - container assignment attempts < 3
scheduler.getProgress();
drainableAppCallback.drain();
verify(mockRMClient, times(0)).releaseAssignedContainer((ContainerId) any());
// no need for task preemption until now - so they should match
Assert.assertEquals(scheduler.numHeartbeats, scheduler.heartbeatAtLastPreemption);
heldContainer.incrementAssignmentAttempts();
// no preemption - container assignment attempts < 3
scheduler.getProgress();
drainableAppCallback.drain();
verify(mockRMClient, times(0)).releaseAssignedContainer((ContainerId) any());
heldContainer.incrementAssignmentAttempts();
heldContainer.incrementAssignmentAttempts();
// preemption - container released and resource asked again
scheduler.getProgress();
drainableAppCallback.drain();
verify(mockRMClient, times(1)).releaseAssignedContainer((ContainerId) any());
verify(mockRMClient, times(1)).releaseAssignedContainer(mockCId4);
// internally re-request pri8 task request because we release pri8 new container
verify(mockRMClient, times(7)).addContainerRequest(requestCaptor.capture());
CookieContainerRequest reAdded = requestCaptor.getValue();
Assert.assertEquals(pri8, reAdded.getPriority());
Assert.assertEquals(taskAsk, reAdded.getCapability());
Assert.assertEquals(mockTaskPri8Cookie, reAdded.getCookie().getAppCookie());
// remove fudging.
scheduler.delayedContainerManager.delayedContainers.clear();
// no need for task preemption until now - so they should match
Assert.assertEquals(scheduler.numHeartbeats, scheduler.heartbeatAtLastPreemption);
scheduler.allocateTask(mockTask3Retry, taskAsk, null, null, pri5, obj3, null);
// no preemption - higher pri. exact match
scheduler.getProgress();
// no need for task preemption until now - so they should match
drainableAppCallback.drain();
// no need for task preemption until now - so they should match
Assert.assertEquals(scheduler.numHeartbeats, scheduler.heartbeatAtLastPreemption);
verify(mockRMClient, times(1)).releaseAssignedContainer((ContainerId) any());
for (int i = 0; i < 11; ++i) {
scheduler.allocateTask(mockTask2, taskAsk, null, null, pri4, null, null);
}
drainableAppCallback.drain();
// mockTaskPri3KillB gets preempted to clear 10% of outstanding running preemptable tasks
// this is also a higher priority container than the pending task priority but was running a
// lower priority task. Task priority is relevant for preemption and not container priority as
// containers can run tasks of different priorities
// first heartbeat
scheduler.getProgress();
Assert.assertTrue(scheduler.numHeartbeats > scheduler.heartbeatAtLastPreemption);
drainableAppCallback.drain();
// second heartbeat
scheduler.getProgress();
drainableAppCallback.drain();
verify(mockRMClient, times(1)).releaseAssignedContainer((ContainerId) any());
// third heartbeat
scheduler.getProgress();
drainableAppCallback.drain();
verify(mockRMClient, times(2)).releaseAssignedContainer((ContainerId) any());
verify(mockRMClient, times(1)).releaseAssignedContainer(mockCId3B);
Assert.assertEquals(scheduler.numHeartbeats, scheduler.heartbeatAtLastPreemption);
// there are pending preemptions.
// first heartbeat
scheduler.getProgress();
// second heartbeat
scheduler.getProgress();
verify(mockRMClient, times(2)).releaseAssignedContainer((ContainerId) any());
// third heartbeat
scheduler.getProgress();
drainableAppCallback.drain();
// Next oldest mockTaskPri3KillA gets preempted to clear 10% of outstanding running preemptable tasks
verify(mockRMClient, times(3)).releaseAssignedContainer((ContainerId) any());
verify(mockRMClient, times(1)).releaseAssignedContainer(mockCId3A);
AppFinalStatus finalStatus = new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, "", DEFAULT_APP_URL);
when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
scheduler.shutdown();
drainableAppCallback.drain();
}
use of org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.TaskSchedulerWithDrainableContext in project tez by apache.
the class TestTaskScheduler method testTaskSchedulerNoReuse.
@SuppressWarnings({ "unchecked" })
@Test(timeout = 10000)
public void testTaskSchedulerNoReuse() throws Exception {
AMRMClientAsyncForTest mockRMClient = spy(new AMRMClientAsyncForTest(new AMRMClientForTest(), 100));
Configuration conf = new Configuration();
conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, false);
int interval = 100;
conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, interval);
TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(DEFAULT_APP_HOST, DEFAULT_APP_PORT, DEFAULT_APP_URL, conf);
TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);
TaskSchedulerWithDrainableContext scheduler = new TaskSchedulerWithDrainableContext(drainableAppCallback, mockRMClient);
scheduler.initialize();
drainableAppCallback.drain();
// Verifying the validity of the configuration via the interval only instead of making sure
// it's the same instance.
verify(mockRMClient).setHeartbeatInterval(interval);
scheduler.start();
drainableAppCallback.drain();
verify(mockRMClient).start();
verify(mockRMClient).registerApplicationMaster(DEFAULT_APP_HOST, DEFAULT_APP_PORT, DEFAULT_APP_URL);
RegisterApplicationMasterResponse regResponse = mockRMClient.getRegistrationResponse();
verify(mockApp).setApplicationRegistrationData(regResponse.getMaximumResourceCapability(), regResponse.getApplicationACLs(), regResponse.getClientToAMTokenMasterKey(), regResponse.getQueue());
Assert.assertEquals(scheduler.getClusterNodeCount(), mockRMClient.getClusterNodeCount());
Object mockTask1 = mock(Object.class);
Object mockCookie1 = mock(Object.class);
Resource mockCapability = mock(Resource.class);
String[] hosts = { "host1", "host5" };
String[] racks = { "/default-rack", "/default-rack" };
Priority mockPriority = mock(Priority.class);
ArgumentCaptor<CookieContainerRequest> requestCaptor = ArgumentCaptor.forClass(CookieContainerRequest.class);
// allocate task
scheduler.allocateTask(mockTask1, mockCapability, hosts, racks, mockPriority, null, mockCookie1);
drainableAppCallback.drain();
verify(mockRMClient, times(1)).addContainerRequest((CookieContainerRequest) any());
// returned from task requests before allocation happens
assertFalse(scheduler.deallocateTask(mockTask1, true, null, null));
verify(mockApp, times(0)).containerBeingReleased(any(ContainerId.class));
verify(mockRMClient, times(1)).removeContainerRequest((CookieContainerRequest) any());
verify(mockRMClient, times(0)).releaseAssignedContainer((ContainerId) any());
// deallocating unknown task
assertFalse(scheduler.deallocateTask(mockTask1, true, null, null));
verify(mockApp, times(0)).containerBeingReleased(any(ContainerId.class));
verify(mockRMClient, times(1)).removeContainerRequest((CookieContainerRequest) any());
verify(mockRMClient, times(0)).releaseAssignedContainer((ContainerId) any());
// allocate tasks
Object mockTask2 = mock(Object.class);
Object mockCookie2 = mock(Object.class);
Object mockTask3 = mock(Object.class);
Object mockCookie3 = mock(Object.class);
scheduler.allocateTask(mockTask1, mockCapability, hosts, racks, mockPriority, null, mockCookie1);
drainableAppCallback.drain();
verify(mockRMClient, times(2)).addContainerRequest(requestCaptor.capture());
CookieContainerRequest request1 = requestCaptor.getValue();
scheduler.allocateTask(mockTask2, mockCapability, hosts, racks, mockPriority, null, mockCookie2);
drainableAppCallback.drain();
verify(mockRMClient, times(3)).addContainerRequest(requestCaptor.capture());
CookieContainerRequest request2 = requestCaptor.getValue();
scheduler.allocateTask(mockTask3, mockCapability, hosts, racks, mockPriority, null, mockCookie3);
drainableAppCallback.drain();
verify(mockRMClient, times(4)).addContainerRequest(requestCaptor.capture());
CookieContainerRequest request3 = requestCaptor.getValue();
List<Container> containers = new ArrayList<Container>();
Container mockContainer1 = mock(Container.class, RETURNS_DEEP_STUBS);
when(mockContainer1.getNodeId().getHost()).thenReturn("host1");
ContainerId mockCId1 = mock(ContainerId.class);
when(mockContainer1.getId()).thenReturn(mockCId1);
containers.add(mockContainer1);
Container mockContainer2 = mock(Container.class, RETURNS_DEEP_STUBS);
when(mockContainer2.getNodeId().getHost()).thenReturn("host2");
ContainerId mockCId2 = mock(ContainerId.class);
when(mockContainer2.getId()).thenReturn(mockCId2);
containers.add(mockContainer2);
Container mockContainer3 = mock(Container.class, RETURNS_DEEP_STUBS);
when(mockContainer3.getNodeId().getHost()).thenReturn("host3");
ContainerId mockCId3 = mock(ContainerId.class);
when(mockContainer3.getId()).thenReturn(mockCId3);
containers.add(mockContainer3);
Container mockContainer4 = mock(Container.class, RETURNS_DEEP_STUBS);
when(mockContainer4.getNodeId().getHost()).thenReturn("host4");
ContainerId mockCId4 = mock(ContainerId.class);
when(mockContainer4.getId()).thenReturn(mockCId4);
containers.add(mockContainer4);
scheduler.onContainersAllocated(containers);
drainableAppCallback.drain();
// first container allocated
verify(mockApp).taskAllocated(mockTask1, mockCookie1, mockContainer1);
verify(mockApp).taskAllocated(mockTask2, mockCookie2, mockContainer2);
verify(mockApp).taskAllocated(mockTask3, mockCookie3, mockContainer3);
// no other allocations returned
verify(mockApp, times(3)).taskAllocated(any(), any(), (Container) any());
verify(mockRMClient).removeContainerRequest(request1);
verify(mockRMClient).removeContainerRequest(request2);
verify(mockRMClient).removeContainerRequest(request3);
// verify unwanted container released
verify(mockRMClient).releaseAssignedContainer(mockCId4);
// deallocate allocated task
assertTrue(scheduler.deallocateTask(mockTask1, true, null, null));
drainableAppCallback.drain();
verify(mockApp).containerBeingReleased(mockCId1);
verify(mockRMClient).releaseAssignedContainer(mockCId1);
// deallocate allocated container
Assert.assertEquals(mockTask2, scheduler.deallocateContainer(mockCId2));
drainableAppCallback.drain();
verify(mockRMClient).releaseAssignedContainer(mockCId2);
verify(mockRMClient, times(3)).releaseAssignedContainer((ContainerId) any());
List<ContainerStatus> statuses = new ArrayList<ContainerStatus>();
ContainerStatus mockStatus1 = mock(ContainerStatus.class);
when(mockStatus1.getContainerId()).thenReturn(mockCId1);
statuses.add(mockStatus1);
ContainerStatus mockStatus2 = mock(ContainerStatus.class);
when(mockStatus2.getContainerId()).thenReturn(mockCId2);
statuses.add(mockStatus2);
ContainerStatus mockStatus3 = mock(ContainerStatus.class);
when(mockStatus3.getContainerId()).thenReturn(mockCId3);
statuses.add(mockStatus3);
ContainerStatus mockStatus4 = mock(ContainerStatus.class);
when(mockStatus4.getContainerId()).thenReturn(mockCId4);
statuses.add(mockStatus4);
scheduler.onContainersCompleted(statuses);
drainableAppCallback.drain();
// released container status returned
verify(mockApp).containerCompleted(mockTask1, mockStatus1);
verify(mockApp).containerCompleted(mockTask2, mockStatus2);
// currently allocated container status returned and not released
verify(mockApp).containerCompleted(mockTask3, mockStatus3);
// no other statuses returned
verify(mockApp, times(3)).containerCompleted(any(), (ContainerStatus) any());
verify(mockRMClient, times(3)).releaseAssignedContainer((ContainerId) any());
// verify blacklisting
verify(mockRMClient, times(0)).addNodeToBlacklist((NodeId) any());
String badHost = "host6";
NodeId badNodeId = mock(NodeId.class);
when(badNodeId.getHost()).thenReturn(badHost);
scheduler.blacklistNode(badNodeId);
verify(mockRMClient, times(1)).addNodeToBlacklist(badNodeId);
Object mockTask4 = mock(Object.class);
Object mockCookie4 = mock(Object.class);
scheduler.allocateTask(mockTask4, mockCapability, null, null, mockPriority, null, mockCookie4);
drainableAppCallback.drain();
verify(mockRMClient, times(5)).addContainerRequest(requestCaptor.capture());
Container mockContainer5 = mock(Container.class, RETURNS_DEEP_STUBS);
when(mockContainer5.getNodeId().getHost()).thenReturn(badHost);
when(mockContainer5.getNodeId()).thenReturn(badNodeId);
ContainerId mockCId5 = mock(ContainerId.class);
when(mockContainer5.getId()).thenReturn(mockCId5);
containers.clear();
containers.add(mockContainer5);
scheduler.onContainersAllocated(containers);
drainableAppCallback.drain();
// no new allocation
verify(mockApp, times(3)).taskAllocated(any(), any(), (Container) any());
// verify blacklisted container released
verify(mockRMClient).releaseAssignedContainer(mockCId5);
verify(mockRMClient, times(4)).releaseAssignedContainer((ContainerId) any());
// verify request added back
verify(mockRMClient, times(6)).addContainerRequest(requestCaptor.capture());
Container mockContainer6 = mock(Container.class, RETURNS_DEEP_STUBS);
when(mockContainer6.getNodeId().getHost()).thenReturn("host7");
ContainerId mockCId6 = mock(ContainerId.class);
when(mockContainer6.getId()).thenReturn(mockCId6);
containers.clear();
containers.add(mockContainer6);
scheduler.onContainersAllocated(containers);
drainableAppCallback.drain();
// new allocation
verify(mockApp, times(4)).taskAllocated(any(), any(), (Container) any());
verify(mockApp).taskAllocated(mockTask4, mockCookie4, mockContainer6);
// deallocate allocated task
assertTrue(scheduler.deallocateTask(mockTask4, true, null, null));
drainableAppCallback.drain();
verify(mockApp).containerBeingReleased(mockCId6);
verify(mockRMClient).releaseAssignedContainer(mockCId6);
verify(mockRMClient, times(5)).releaseAssignedContainer((ContainerId) any());
// test unblacklist
scheduler.unblacklistNode(badNodeId);
verify(mockRMClient, times(1)).removeNodeFromBlacklist(badNodeId);
assertEquals(0, scheduler.blacklistedNodes.size());
float progress = 0.5f;
when(mockApp.getProgress()).thenReturn(progress);
Assert.assertEquals(progress, scheduler.getProgress(), 0);
// check duplicate allocation request
scheduler.allocateTask(mockTask1, mockCapability, hosts, racks, mockPriority, null, mockCookie1);
drainableAppCallback.drain();
verify(mockRMClient, times(7)).addContainerRequest((CookieContainerRequest) any());
verify(mockRMClient, times(6)).removeContainerRequest((CookieContainerRequest) any());
scheduler.allocateTask(mockTask1, mockCapability, hosts, racks, mockPriority, null, mockCookie1);
drainableAppCallback.drain();
// old request removed and new one added
verify(mockRMClient, times(7)).removeContainerRequest((CookieContainerRequest) any());
verify(mockRMClient, times(8)).addContainerRequest((CookieContainerRequest) any());
assertFalse(scheduler.deallocateTask(mockTask1, true, null, null));
List<NodeReport> mockUpdatedNodes = mock(List.class);
scheduler.onNodesUpdated(mockUpdatedNodes);
drainableAppCallback.drain();
verify(mockApp).nodesUpdated(mockUpdatedNodes);
ArgumentCaptor<String> argumentCaptor = ArgumentCaptor.forClass(String.class);
Exception mockException = new IOException("mockexception");
scheduler.onError(mockException);
drainableAppCallback.drain();
verify(mockApp).reportError(eq(YarnTaskSchedulerServiceError.RESOURCEMANAGER_ERROR), argumentCaptor.capture(), any(DagInfo.class));
assertTrue(argumentCaptor.getValue().contains("mockexception"));
scheduler.onShutdownRequest();
drainableAppCallback.drain();
verify(mockApp).appShutdownRequested();
AppFinalStatus finalStatus = new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, SUCCEED_APP_MESSAGE, DEFAULT_APP_URL);
when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
scheduler.shutdown();
drainableAppCallback.drain();
verify(mockRMClient).unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED, SUCCEED_APP_MESSAGE, DEFAULT_APP_URL);
verify(mockRMClient).stop();
}
Aggregations