@Test(timeout = 10000)
public void testDAGSchedulerNaturalOrder() {
DAG mockDag = mock(DAG.class);
Vertex mockVertex = mock(Vertex.class);
TaskAttempt mockAttempt = mock(TaskAttempt.class);
when(mockDag.getVertex((TezVertexID) any())).thenReturn(mockVertex);
when(mockVertex.getDistanceFromRoot()).thenReturn(0).thenReturn(1).thenReturn(2);
when(mockAttempt.getIsRescheduled()).thenReturn(false);
DAGEventSchedulerUpdate event =
new DAGEventSchedulerUpdate(DAGEventSchedulerUpdate.UpdateType.TA_SCHEDULE, mockAttempt);
DAGScheduler scheduler = new DAGSchedulerNaturalOrder(mockDag, mockEventHandler);
scheduler.scheduleTask(event);
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 2);
scheduler.scheduleTask(event);
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 4);
scheduler.scheduleTask(event);
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 6);
when(mockAttempt.getIsRescheduled()).thenReturn(true);
scheduler.scheduleTask(event);
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 5);
}
@Ignore
@Test(timeout = 10000)
public void testDAGSchedulerMRR() {
DAG mockDag = mock(DAG.class);
TezDAGID dagId = TezDAGID.getInstance("1", 1, 1);
TaskSchedulerEventHandler mockTaskScheduler = mock(TaskSchedulerEventHandler.class);
Vertex mockVertex1 = mock(Vertex.class);
TezVertexID mockVertexId1 = TezVertexID.getInstance(dagId, 1);
when(mockVertex1.getVertexId()).thenReturn(mockVertexId1);
when(mockVertex1.getDistanceFromRoot()).thenReturn(0);
TaskAttempt mockAttempt1 = mock(TaskAttempt.class);
when(mockAttempt1.getVertexID()).thenReturn(mockVertexId1);
when(mockAttempt1.getIsRescheduled()).thenReturn(false);
when(mockDag.getVertex(mockVertexId1)).thenReturn(mockVertex1);
Vertex mockVertex2 = mock(Vertex.class);
TezVertexID mockVertexId2 = TezVertexID.getInstance(dagId, 2);
when(mockVertex2.getVertexId()).thenReturn(mockVertexId2);
when(mockVertex2.getDistanceFromRoot()).thenReturn(1);
TaskAttempt mockAttempt2 = mock(TaskAttempt.class);
when(mockAttempt2.getVertexID()).thenReturn(mockVertexId2);
when(mockAttempt2.getIsRescheduled()).thenReturn(false);
when(mockDag.getVertex(mockVertexId2)).thenReturn(mockVertex2);
TaskAttempt mockAttempt2f = mock(TaskAttempt.class);
when(mockAttempt2f.getVertexID()).thenReturn(mockVertexId2);
when(mockAttempt2f.getIsRescheduled()).thenReturn(true);
Vertex mockVertex3 = mock(Vertex.class);
TezVertexID mockVertexId3 = TezVertexID.getInstance(dagId, 3);
when(mockVertex3.getVertexId()).thenReturn(mockVertexId3);
when(mockVertex3.getDistanceFromRoot()).thenReturn(2);
TaskAttempt mockAttempt3 = mock(TaskAttempt.class);
when(mockAttempt3.getVertexID()).thenReturn(mockVertexId3);
when(mockAttempt3.getIsRescheduled()).thenReturn(false);
when(mockDag.getVertex(mockVertexId3)).thenReturn(mockVertex3);
DAGEventSchedulerUpdate mockEvent1 = mock(DAGEventSchedulerUpdate.class);
when(mockEvent1.getAttempt()).thenReturn(mockAttempt1);
DAGEventSchedulerUpdate mockEvent2 = mock(DAGEventSchedulerUpdate.class);
when(mockEvent2.getAttempt()).thenReturn(mockAttempt2);
DAGEventSchedulerUpdate mockEvent2f = mock(DAGEventSchedulerUpdate.class);
when(mockEvent2f.getAttempt()).thenReturn(mockAttempt2f);
DAGEventSchedulerUpdate mockEvent3 = mock(DAGEventSchedulerUpdate.class);
when(mockEvent3.getAttempt()).thenReturn(mockAttempt3);
DAGScheduler scheduler =
new DAGSchedulerMRR(mockDag, mockEventHandler, mockTaskScheduler, 0.5f);
// M starts. M completes. R1 starts. R1 completes. R2 starts. R2 completes
scheduler.scheduleTask(mockEvent1); // M starts
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 3);
scheduler.scheduleTask(mockEvent1); // M runs another
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 3);
scheduler.vertexCompleted(mockVertex1); // M completes
scheduler.scheduleTask(mockEvent2); // R1 starts
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 6);
scheduler.scheduleTask(mockEvent2); // R1 runs another
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 6);
scheduler.scheduleTask(mockEvent2f); // R1 runs retry. Retry priority
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 4);
scheduler.vertexCompleted(mockVertex2); // R1 completes
scheduler.scheduleTask(mockEvent3); // R2 starts
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 9);
scheduler.scheduleTask(mockEvent3); // R2 runs another
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 9);
scheduler.vertexCompleted(mockVertex3); // R2 completes
// M starts. R1 starts. M completes. R2 starts. R1 completes. R2 completes
scheduler.scheduleTask(mockEvent1); // M starts
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 3);
scheduler.scheduleTask(mockEvent2); // R1 starts. Reordered priority
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 2);
scheduler.scheduleTask(mockEvent1); // M runs another
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 3);
scheduler.scheduleTask(mockEvent2); // R1 runs another. Reordered priority
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 2);
scheduler.scheduleTask(mockEvent2f); // R1 runs retry. Reordered priority
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 2);
scheduler.vertexCompleted(mockVertex1); // M completes
scheduler.scheduleTask(mockEvent3); // R2 starts. Reordered priority
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 5);
scheduler.scheduleTask(mockEvent2); // R1 runs another. Normal priority
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 6);
scheduler.scheduleTask(mockEvent2f); // R1 runs retry. Retry priority
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 4);
scheduler.scheduleTask(mockEvent3); // R2 runs another. Reordered priority
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 5);
scheduler.vertexCompleted(mockVertex2); // R1 completes
scheduler.vertexCompleted(mockVertex3); // R2 completes
// M starts. M completes. R1 starts. R2 starts. R1 completes. R2 completes
scheduler.scheduleTask(mockEvent1); // M starts
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 3);
scheduler.vertexCompleted(mockVertex1); // M completes
scheduler.scheduleTask(mockEvent2); // R1 starts
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 6);
scheduler.scheduleTask(mockEvent3); // R2 starts. Reordered priority
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 5);
scheduler.scheduleTask(mockEvent2); // R1 runs another
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 6);
scheduler.vertexCompleted(mockVertex2); // R1 completes
scheduler.vertexCompleted(mockVertex3); // R2 completes
// M starts. R1 starts. M completes. R1 completes. R2 starts. R2 completes
scheduler.scheduleTask(mockEvent1); // M starts
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 3);
scheduler.scheduleTask(mockEvent2); // R1 starts. Reordered priority
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 2);
scheduler.vertexCompleted(mockVertex1); // M completes
scheduler.scheduleTask(mockEvent2); // R1 starts. Normal priority
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 6);
scheduler.vertexCompleted(mockVertex2); // R1 completes
scheduler.scheduleTask(mockEvent3); // R2 starts
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 9);
scheduler.vertexCompleted(mockVertex3); // R2 completes
}