Examples with EdgeProperty org.apache.tez.dag.api.EdgeProperty used on opensource projects

Example 51 with EdgeProperty

use of org.apache.tez.dag.api.EdgeProperty in project tez by apache.

the class TestInputReadyVertexManager method testComplex.

@Test(timeout = 5000)
public void testComplex() throws Exception {
    HashMap<String, EdgeProperty> mockInputVertices = new HashMap<String, EdgeProperty>();
    String mockSrcVertexId1 = "Vertex1";
    EdgeProperty eProp1 = EdgeProperty.create(EdgeProperty.DataMovementType.SCATTER_GATHER, EdgeProperty.DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, OutputDescriptor.create("out"), InputDescriptor.create("in"));
    String mockSrcVertexId2 = "Vertex2";
    EdgeProperty eProp2 = EdgeProperty.create(EdgeProperty.DataMovementType.ONE_TO_ONE, EdgeProperty.DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, OutputDescriptor.create("out"), InputDescriptor.create("in"));
    String mockSrcVertexId3 = "Vertex3";
    EdgeProperty eProp3 = EdgeProperty.create(EdgeProperty.DataMovementType.ONE_TO_ONE, EdgeProperty.DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, OutputDescriptor.create("out"), InputDescriptor.create("in"));
    String mockManagedVertexId = "Vertex";
    Container mockContainer2 = mock(Container.class);
    ContainerId mockCId2 = mock(ContainerId.class);
    when(mockContainer2.getId()).thenReturn(mockCId2);
    Container mockContainer3 = mock(Container.class);
    ContainerId mockCId3 = mock(ContainerId.class);
    when(mockContainer3.getId()).thenReturn(mockCId3);
    VertexManagerPluginContext mockContext = mock(VertexManagerPluginContext.class);
    when(mockContext.getInputVertexEdgeProperties()).thenReturn(mockInputVertices);
    when(mockContext.getVertexName()).thenReturn(mockManagedVertexId);
    when(mockContext.getVertexNumTasks(mockSrcVertexId1)).thenReturn(3);
    when(mockContext.getVertexNumTasks(mockSrcVertexId2)).thenReturn(3);
    when(mockContext.getVertexNumTasks(mockSrcVertexId3)).thenReturn(3);
    mockInputVertices.put(mockSrcVertexId1, eProp1);
    mockInputVertices.put(mockSrcVertexId2, eProp2);
    mockInputVertices.put(mockSrcVertexId3, eProp3);
    List<TaskAttemptIdentifier> initialCompletions = Lists.newArrayList();
    // 1-1 sources do not match managed tasks. setParallelism called to make them match
    when(mockContext.getVertexNumTasks(mockManagedVertexId)).thenReturn(4);
    InputReadyVertexManager manager = new InputReadyVertexManager(mockContext);
    manager.initialize();
    verify(mockContext, times(1)).vertexReconfigurationPlanned();
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId1, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId2, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId3, VertexState.CONFIGURED));
    verify(mockContext, times(1)).reconfigureVertex(3, null, null);
    verify(mockContext, times(1)).doneReconfiguringVertex();
    manager.onVertexStarted(initialCompletions);
    // 1-1 sources do not match
    when(mockContext.getVertexNumTasks(mockManagedVertexId)).thenReturn(3);
    when(mockContext.getVertexNumTasks(mockSrcVertexId3)).thenReturn(4);
    manager = new InputReadyVertexManager(mockContext);
    manager.initialize();
    verify(mockContext, times(2)).vertexReconfigurationPlanned();
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId1, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId2, VertexState.CONFIGURED));
    try {
        manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId3, VertexState.CONFIGURED));
        Assert.assertTrue("Should have exception", false);
    } catch (TezUncheckedException e) {
        e.getMessage().contains("1-1 source vertices must have identical concurrency");
    }
    verify(mockContext, times(1)).reconfigureVertex(anyInt(), (VertexLocationHint) any(), // not invoked
    anyMap());
    when(mockContext.getVertexNumTasks(mockSrcVertexId3)).thenReturn(3);
    initialCompletions.add(TestShuffleVertexManager.createTaskAttemptIdentifier(mockSrcVertexId1, 0));
    initialCompletions.add(TestShuffleVertexManager.createTaskAttemptIdentifier(mockSrcVertexId2, 0));
    manager = new InputReadyVertexManager(mockContext);
    manager.initialize();
    verify(mockContext, times(3)).vertexReconfigurationPlanned();
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId1, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId2, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId3, VertexState.CONFIGURED));
    verify(mockContext, times(1)).reconfigureVertex(anyInt(), (VertexLocationHint) any(), // not invoked
    anyMap());
    verify(mockContext, times(2)).doneReconfiguringVertex();
    manager.onVertexStarted(initialCompletions);
    // all 1-1 0's done but not scheduled because v1 is not done
    manager.onSourceTaskCompleted(TestShuffleVertexManager.createTaskAttemptIdentifier(mockSrcVertexId3, 0));
    manager.onSourceTaskCompleted(TestShuffleVertexManager.createTaskAttemptIdentifier(mockSrcVertexId1, 1));
    manager.onSourceTaskCompleted(// duplicate
    TestShuffleVertexManager.createTaskAttemptIdentifier(mockSrcVertexId1, 1));
    manager.onSourceTaskCompleted(TestShuffleVertexManager.createTaskAttemptIdentifier(mockSrcVertexId2, 1));
    verify(mockContext, times(0)).scheduleTasks(anyList());
    manager.onSourceTaskCompleted(// v1 done
    TestShuffleVertexManager.createTaskAttemptIdentifier(mockSrcVertexId1, 2));
    verify(mockContext, times(1)).scheduleTasks(requestCaptor.capture());
    Assert.assertEquals(1, requestCaptor.getValue().size());
    Assert.assertEquals(0, requestCaptor.getValue().get(0).getTaskIndex());
    Assert.assertEquals(mockSrcVertexId3, requestCaptor.getValue().get(0).getTaskLocationHint().getAffinitizedTask().getVertexName());
    Assert.assertEquals(0, requestCaptor.getValue().get(0).getTaskLocationHint().getAffinitizedTask().getTaskIndex());
    // 1-1 completion triggers since other 1-1 is done
    manager.onSourceTaskCompleted(TestShuffleVertexManager.createTaskAttemptIdentifier(mockSrcVertexId3, 1));
    verify(mockContext, times(2)).scheduleTasks(requestCaptor.capture());
    Assert.assertEquals(1, requestCaptor.getValue().size());
    Assert.assertEquals(1, requestCaptor.getValue().get(0).getTaskIndex());
    Assert.assertEquals(mockSrcVertexId3, requestCaptor.getValue().get(0).getTaskLocationHint().getAffinitizedTask().getVertexName());
    Assert.assertEquals(1, requestCaptor.getValue().get(0).getTaskLocationHint().getAffinitizedTask().getTaskIndex());
    // 1-1 completion does not trigger since other 1-1 is not done
    manager.onSourceTaskCompleted(TestShuffleVertexManager.createTaskAttemptIdentifier(mockSrcVertexId3, 2));
    verify(mockContext, times(2)).scheduleTasks(anyList());
    // 1-1 completion trigger start
    manager.onSourceTaskCompleted(TestShuffleVertexManager.createTaskAttemptIdentifier(mockSrcVertexId2, 2));
    verify(mockContext, times(3)).scheduleTasks(requestCaptor.capture());
    Assert.assertEquals(1, requestCaptor.getValue().size());
    Assert.assertEquals(2, requestCaptor.getValue().get(0).getTaskIndex());
    Assert.assertEquals(mockSrcVertexId2, requestCaptor.getValue().get(0).getTaskLocationHint().getAffinitizedTask().getVertexName());
    Assert.assertEquals(2, requestCaptor.getValue().get(0).getTaskLocationHint().getAffinitizedTask().getTaskIndex());
    // no more starts
    manager.onSourceTaskCompleted(TestShuffleVertexManager.createTaskAttemptIdentifier(mockSrcVertexId3, 2));
    verify(mockContext, times(3)).scheduleTasks(anyList());
}

Example 52 with EdgeProperty

use of org.apache.tez.dag.api.EdgeProperty in project tez by apache.

the class TestShuffleVertexManagerBase method test_Tez1649_with_scatter_gather_edges.

/**
 * Tasks should be scheduled only when all source vertices are configured completely
 * @throws IOException
 */
@Test(timeout = 5000)
public void test_Tez1649_with_scatter_gather_edges() throws IOException {
    Configuration conf = new Configuration();
    ShuffleVertexManagerBase manager = null;
    HashMap<String, EdgeProperty> mockInputVertices_R2 = new HashMap<String, EdgeProperty>();
    String r1 = "R1";
    EdgeProperty eProp1 = EdgeProperty.create(EdgeProperty.DataMovementType.SCATTER_GATHER, EdgeProperty.DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, OutputDescriptor.create("out"), InputDescriptor.create("in"));
    String m2 = "M2";
    EdgeProperty eProp2 = EdgeProperty.create(EdgeProperty.DataMovementType.SCATTER_GATHER, EdgeProperty.DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, OutputDescriptor.create("out"), InputDescriptor.create("in"));
    String m3 = "M3";
    EdgeProperty eProp3 = EdgeProperty.create(EdgeProperty.DataMovementType.SCATTER_GATHER, EdgeProperty.DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, OutputDescriptor.create("out"), InputDescriptor.create("in"));
    final String mockManagedVertexId_R2 = "R2";
    mockInputVertices_R2.put(r1, eProp1);
    mockInputVertices_R2.put(m2, eProp2);
    mockInputVertices_R2.put(m3, eProp3);
    final VertexManagerPluginContext mockContext_R2 = mock(VertexManagerPluginContext.class);
    when(mockContext_R2.getInputVertexEdgeProperties()).thenReturn(mockInputVertices_R2);
    when(mockContext_R2.getVertexName()).thenReturn(mockManagedVertexId_R2);
    when(mockContext_R2.getVertexNumTasks(mockManagedVertexId_R2)).thenReturn(3);
    when(mockContext_R2.getVertexNumTasks(r1)).thenReturn(3);
    when(mockContext_R2.getVertexNumTasks(m2)).thenReturn(3);
    when(mockContext_R2.getVertexNumTasks(m3)).thenReturn(3);
    VertexManagerEvent vmEvent = getVertexManagerEvent(null, 50L, r1);
    // check initialization
    manager = createManager(conf, mockContext_R2, 0.001f, 0.001f);
    final List<Integer> scheduledTasks = Lists.newLinkedList();
    doAnswer(new ScheduledTasksAnswer(scheduledTasks)).when(mockContext_R2).scheduleTasks(anyList());
    manager.onVertexStarted(emptyCompletions);
    Assert.assertTrue(manager.bipartiteSources == 3);
    manager.onVertexStateUpdated(new VertexStateUpdate(m2, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(m3, VertexState.CONFIGURED));
    manager.onVertexManagerEventReceived(vmEvent);
    // no tasks scheduled
    Assert.assertEquals(3, manager.pendingTasks.size());
    Assert.assertEquals(6, manager.totalNumBipartiteSourceTasks);
    Assert.assertEquals(0, manager.numBipartiteSourceTasksCompleted);
    // no tasks scheduled
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    Assert.assertTrue(manager.totalNumBipartiteSourceTasks == 6);
    // Send events for all tasks of m3.
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(m3, 0));
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(m3, 1));
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(m3, 2));
    // no tasks scheduled
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    Assert.assertTrue(manager.totalNumBipartiteSourceTasks == 6);
    // Send events for m2. But still we need to wait for at least 1 event from r1.
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(m2, 0));
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(m2, 1));
    // no tasks scheduled
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    Assert.assertTrue(manager.totalNumBipartiteSourceTasks == 6);
    // we need to wait for at least 1 event from r1 to make sure all vertices cross min threshold
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(r1, 0));
    // no tasks scheduled
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    Assert.assertTrue(manager.totalNumBipartiteSourceTasks == 6);
    // Ensure that setVertexParallelism is not called for R2.
    verify(mockContext_R2, times(0)).reconfigureVertex(anyInt(), any(VertexLocationHint.class), anyMap());
    // ShuffleVertexManager's updatePendingTasks relies on getVertexNumTasks. Setting this for test
    when(mockContext_R2.getVertexNumTasks(mockManagedVertexId_R2)).thenReturn(1);
    // complete configuration of r1 triggers the scheduling
    manager.onVertexStateUpdated(new VertexStateUpdate(r1, VertexState.CONFIGURED));
    Assert.assertTrue(manager.totalNumBipartiteSourceTasks == 9);
    verify(mockContext_R2, times(1)).reconfigureVertex(eq(1), any(VertexLocationHint.class), anyMap());
    // all tasks scheduled
    Assert.assertTrue(manager.pendingTasks.size() == 0);
    Assert.assertTrue(scheduledTasks.size() == 1);
    // try with zero task vertices
    scheduledTasks.clear();
    when(mockContext_R2.getInputVertexEdgeProperties()).thenReturn(mockInputVertices_R2);
    when(mockContext_R2.getVertexName()).thenReturn(mockManagedVertexId_R2);
    when(mockContext_R2.getVertexNumTasks(mockManagedVertexId_R2)).thenReturn(3);
    when(mockContext_R2.getVertexNumTasks(r1)).thenReturn(0);
    when(mockContext_R2.getVertexNumTasks(m2)).thenReturn(0);
    when(mockContext_R2.getVertexNumTasks(m3)).thenReturn(3);
    manager = createManager(conf, mockContext_R2, 0.001f, 0.001f);
    manager.onVertexStarted(emptyCompletions);
    // no tasks scheduled
    Assert.assertEquals(3, manager.pendingTasks.size());
    Assert.assertEquals(0, manager.numBipartiteSourceTasksCompleted);
    // Only need completed configuration notification from m3
    manager.onVertexStateUpdated(new VertexStateUpdate(m3, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(m2, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(r1, VertexState.CONFIGURED));
    Assert.assertEquals(3, manager.totalNumBipartiteSourceTasks);
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(m3, 0));
    // all tasks scheduled
    Assert.assertTrue(manager.pendingTasks.size() == 0);
    Assert.assertTrue(scheduledTasks.size() == 3);
}

Example 53 with EdgeProperty

use of org.apache.tez.dag.api.EdgeProperty in project tez by apache.

the class TestShuffleVertexManagerBase method testShuffleVertexManagerSlowStart.

@Test(timeout = 5000)
public void testShuffleVertexManagerSlowStart() {
    Configuration conf = new Configuration();
    ShuffleVertexManagerBase manager = null;
    HashMap<String, EdgeProperty> mockInputVertices = new HashMap<String, EdgeProperty>();
    String mockSrcVertexId1 = "Vertex1";
    EdgeProperty eProp1 = EdgeProperty.create(EdgeProperty.DataMovementType.SCATTER_GATHER, EdgeProperty.DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, OutputDescriptor.create("out"), InputDescriptor.create("in"));
    String mockSrcVertexId2 = "Vertex2";
    EdgeProperty eProp2 = EdgeProperty.create(EdgeProperty.DataMovementType.SCATTER_GATHER, EdgeProperty.DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, OutputDescriptor.create("out"), InputDescriptor.create("in"));
    String mockSrcVertexId3 = "Vertex3";
    EdgeProperty eProp3 = EdgeProperty.create(EdgeProperty.DataMovementType.BROADCAST, EdgeProperty.DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, OutputDescriptor.create("out"), InputDescriptor.create("in"));
    String mockManagedVertexId = "Vertex4";
    VertexManagerPluginContext mockContext = mock(VertexManagerPluginContext.class);
    when(mockContext.getVertexStatistics(any(String.class))).thenReturn(mock(VertexStatistics.class));
    when(mockContext.getInputVertexEdgeProperties()).thenReturn(mockInputVertices);
    when(mockContext.getVertexName()).thenReturn(mockManagedVertexId);
    when(mockContext.getVertexNumTasks(mockManagedVertexId)).thenReturn(3);
    when(mockContext.getVertexNumTasks(mockSrcVertexId3)).thenReturn(1);
    // fail if there is no bipartite src vertex
    mockInputVertices.put(mockSrcVertexId3, eProp3);
    try {
        manager = createManager(conf, mockContext, 0.1f, 0.1f);
        manager.onVertexStarted(emptyCompletions);
        Assert.assertFalse(true);
    } catch (TezUncheckedException e) {
        Assert.assertTrue(e.getMessage().contains("At least 1 bipartite source should exist"));
    }
    mockInputVertices.put(mockSrcVertexId1, eProp1);
    mockInputVertices.put(mockSrcVertexId2, eProp2);
    // check initialization
    manager = createManager(conf, mockContext, 0.1f, 0.1f);
    manager.onVertexStarted(emptyCompletions);
    Assert.assertTrue(manager.bipartiteSources == 2);
    final List<Integer> scheduledTasks = Lists.newLinkedList();
    doAnswer(new ScheduledTasksAnswer(scheduledTasks)).when(mockContext).scheduleTasks(anyList());
    // source vertices have 0 tasks. immediate start of all managed tasks
    when(mockContext.getVertexNumTasks(mockSrcVertexId1)).thenReturn(0);
    when(mockContext.getVertexNumTasks(mockSrcVertexId2)).thenReturn(0);
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId1, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId2, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId3, VertexState.CONFIGURED));
    manager.onVertexStarted(emptyCompletions);
    Assert.assertTrue(manager.pendingTasks.isEmpty());
    // all tasks scheduled
    Assert.assertTrue(scheduledTasks.size() == 3);
    when(mockContext.getVertexNumTasks(mockSrcVertexId1)).thenReturn(2);
    when(mockContext.getVertexNumTasks(mockSrcVertexId2)).thenReturn(2);
    try {
        // source vertex have some tasks. min < 0.
        manager = createManager(conf, mockContext, -0.1f, 0.0f);
        // should not come here
        Assert.assertTrue(false);
    } catch (IllegalArgumentException e) {
        Assert.assertTrue(e.getMessage().contains("Invalid values for slowStartMinFraction"));
    }
    try {
        // source vertex have some tasks. max > 1.
        manager = createManager(conf, mockContext, 0.0f, 95.0f);
        // should not come here
        Assert.assertTrue(false);
    } catch (IllegalArgumentException e) {
        Assert.assertTrue(e.getMessage().contains("Invalid values for slowStartMinFraction"));
    }
    try {
        // source vertex have some tasks. min > max
        manager = createManager(conf, mockContext, 0.5f, 0.3f);
        // should not come here
        Assert.assertTrue(false);
    } catch (IllegalArgumentException e) {
        Assert.assertTrue(e.getMessage().contains("Invalid values for slowStartMinFraction"));
    }
    // source vertex have some tasks. min > default and max undefined
    int numTasks = 20;
    when(mockContext.getVertexNumTasks(mockSrcVertexId1)).thenReturn(numTasks);
    when(mockContext.getVertexNumTasks(mockSrcVertexId2)).thenReturn(numTasks);
    scheduledTasks.clear();
    manager = createManager(conf, mockContext, 0.8f, null);
    manager.onVertexStarted(emptyCompletions);
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId1, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId2, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId3, VertexState.CONFIGURED));
    Assert.assertEquals(3, manager.pendingTasks.size());
    Assert.assertEquals(numTasks * 2, manager.totalNumBipartiteSourceTasks);
    Assert.assertEquals(0, manager.numBipartiteSourceTasksCompleted);
    float completedTasksThreshold = 0.8f * numTasks;
    // Finish all tasks before exceeding the threshold
    for (String mockSrcVertex : new String[] { mockSrcVertexId1, mockSrcVertexId2 }) {
        for (int i = 0; i < mockContext.getVertexNumTasks(mockSrcVertex); ++i) {
            // complete 0th tasks outside the loop
            manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertex, i + 1));
            if ((i + 2) >= completedTasksThreshold) {
                // stop before completing more than min/max source tasks
                break;
            }
        }
    }
    // Since we haven't exceeded the threshold, all tasks are still pending
    Assert.assertEquals(manager.totalTasksToSchedule, manager.pendingTasks.size());
    // no tasks scheduled
    Assert.assertEquals(0, scheduledTasks.size());
    // Cross the threshold min/max threshold to schedule all tasks
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId1, 0));
    Assert.assertEquals(3, manager.pendingTasks.size());
    Assert.assertEquals(0, scheduledTasks.size());
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId2, 0));
    Assert.assertEquals(0, manager.pendingTasks.size());
    // all tasks scheduled
    Assert.assertEquals(manager.totalTasksToSchedule, scheduledTasks.size());
    // reset vertices for next test
    when(mockContext.getVertexNumTasks(mockSrcVertexId1)).thenReturn(2);
    when(mockContext.getVertexNumTasks(mockSrcVertexId2)).thenReturn(2);
    // source vertex have some tasks. min, max == 0
    manager = createManager(conf, mockContext, 0.0f, 0.0f);
    manager.onVertexStarted(emptyCompletions);
    Assert.assertTrue(manager.totalTasksToSchedule == 3);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 0);
    // all source vertices need to be configured
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId1, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId2, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId3, VertexState.CONFIGURED));
    Assert.assertTrue(manager.totalNumBipartiteSourceTasks == 4);
    Assert.assertTrue(manager.pendingTasks.isEmpty());
    // all tasks scheduled
    Assert.assertTrue(scheduledTasks.size() == 3);
    // min, max > 0 and min == max
    manager = createManager(conf, mockContext, 0.25f, 0.25f);
    manager.onVertexStarted(emptyCompletions);
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId1, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId2, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId3, VertexState.CONFIGURED));
    // no tasks scheduled
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    Assert.assertTrue(manager.totalNumBipartiteSourceTasks == 4);
    // task completion from non-bipartite stage does nothing
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId3, 0));
    // no tasks scheduled
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    Assert.assertTrue(manager.totalNumBipartiteSourceTasks == 4);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 0);
    // task completion on only 1 SG edge does nothing
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId2, 0));
    // no tasks scheduled
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    Assert.assertTrue(manager.totalNumBipartiteSourceTasks == 4);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 1);
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId1, 0));
    Assert.assertTrue(manager.pendingTasks.isEmpty());
    // all tasks scheduled
    Assert.assertTrue(scheduledTasks.size() == 3);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 2);
    // min, max > 0 and min == max == absolute max 1.0
    manager = createManager(conf, mockContext, 1.0f, 1.0f);
    manager.onVertexStarted(emptyCompletions);
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId1, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId2, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId3, VertexState.CONFIGURED));
    // no tasks scheduled
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    Assert.assertTrue(manager.totalNumBipartiteSourceTasks == 4);
    // task completion from non-bipartite stage does nothing
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId3, 0));
    // no tasks scheduled
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    Assert.assertTrue(manager.totalNumBipartiteSourceTasks == 4);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 0);
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId1, 0));
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 1);
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId1, 1));
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 2);
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId2, 0));
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 3);
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId2, 1));
    Assert.assertTrue(manager.pendingTasks.isEmpty());
    // all tasks scheduled
    Assert.assertTrue(scheduledTasks.size() == 3);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 4);
    // min, max > 0 and min == max
    manager = createManager(conf, mockContext, 1.0f, 1.0f);
    manager.onVertexStarted(emptyCompletions);
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId1, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId2, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId3, VertexState.CONFIGURED));
    // no tasks scheduled
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    Assert.assertTrue(manager.totalNumBipartiteSourceTasks == 4);
    // task completion from non-bipartite stage does nothing
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId3, 0));
    // no tasks scheduled
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    Assert.assertTrue(manager.totalNumBipartiteSourceTasks == 4);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 0);
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId1, 0));
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 1);
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId1, 1));
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 2);
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId2, 0));
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 3);
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId2, 1));
    Assert.assertTrue(manager.pendingTasks.isEmpty());
    // all tasks scheduled
    Assert.assertTrue(scheduledTasks.size() == 3);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 4);
    // reset vertices for next test
    when(mockContext.getVertexNumTasks(mockSrcVertexId1)).thenReturn(4);
    when(mockContext.getVertexNumTasks(mockSrcVertexId2)).thenReturn(4);
    // min, max > and min < max
    manager = createManager(conf, mockContext, 0.25f, 0.75f);
    manager.onVertexStarted(emptyCompletions);
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId1, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId2, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId3, VertexState.CONFIGURED));
    // no tasks scheduled
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    Assert.assertTrue(manager.totalNumBipartiteSourceTasks == 8);
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId1, 0));
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId2, 1));
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 2);
    // completion of same task again should not get counted
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId2, 1));
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 2);
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId1, 1));
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId2, 0));
    Assert.assertTrue(manager.pendingTasks.size() == 1);
    // 2 task scheduled
    Assert.assertTrue(scheduledTasks.size() == 2);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 4);
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId1, 2));
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId2, 2));
    Assert.assertTrue(manager.pendingTasks.size() == 0);
    // 1 tasks scheduled
    Assert.assertTrue(scheduledTasks.size() == 1);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 6);
    scheduledTasks.clear();
    // we are done. no action
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId1, 3));
    Assert.assertTrue(manager.pendingTasks.size() == 0);
    // no task scheduled
    Assert.assertTrue(scheduledTasks.size() == 0);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 7);
    // min, max > and min < max
    manager = createManager(conf, mockContext, 0.25f, 1.0f);
    manager.onVertexStarted(emptyCompletions);
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId1, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId2, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId3, VertexState.CONFIGURED));
    // no tasks scheduled
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    Assert.assertTrue(manager.totalNumBipartiteSourceTasks == 8);
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId1, 0));
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId2, 1));
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId1, 1));
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId2, 0));
    Assert.assertTrue(manager.pendingTasks.size() == 2);
    // 1 task scheduled
    Assert.assertTrue(scheduledTasks.size() == 1);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 4);
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId1, 2));
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId2, 2));
    Assert.assertTrue(manager.pendingTasks.size() == 1);
    // 1 task scheduled
    Assert.assertTrue(scheduledTasks.size() == 1);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 6);
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId1, 3));
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId2, 3));
    Assert.assertTrue(manager.pendingTasks.size() == 0);
    // no task scheduled
    Assert.assertTrue(scheduledTasks.size() == 1);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 8);
    // if there is single task to schedule, it should be schedule when src completed
    // fraction is more than min slow start fraction
    scheduledTasks.clear();
    when(mockContext.getVertexNumTasks(mockManagedVertexId)).thenReturn(1);
    manager = createManager(conf, mockContext, 0.25f, 0.75f);
    manager.onVertexStarted(emptyCompletions);
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId1, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId2, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId3, VertexState.CONFIGURED));
    // no tasks scheduled
    Assert.assertTrue(manager.pendingTasks.size() == 1);
    Assert.assertTrue(manager.totalNumBipartiteSourceTasks == 8);
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId1, 0));
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId2, 1));
    Assert.assertTrue(manager.pendingTasks.size() == 1);
    // no task scheduled
    Assert.assertTrue(scheduledTasks.size() == 0);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 2);
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId1, 1));
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId2, 0));
    Assert.assertTrue(manager.pendingTasks.size() == 0);
    // 1 task scheduled
    Assert.assertTrue(scheduledTasks.size() == 1);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 4);
    scheduledTasks.clear();
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId1, 2));
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId2, 2));
    Assert.assertTrue(manager.pendingTasks.size() == 0);
    // no task scheduled
    Assert.assertTrue(scheduledTasks.size() == 0);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 6);
    scheduledTasks.clear();
    // we are done. no action
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(mockSrcVertexId1, 3));
    Assert.assertTrue(manager.pendingTasks.size() == 0);
    // no task scheduled
    Assert.assertTrue(scheduledTasks.size() == 0);
    Assert.assertTrue(manager.numBipartiteSourceTasksCompleted == 7);
}

Example 54 with EdgeProperty

use of org.apache.tez.dag.api.EdgeProperty in project tez by apache.

the class TestShuffleVertexManagerBase method testZeroTasksSendsConfigured.

@Test
public void testZeroTasksSendsConfigured() throws IOException {
    Configuration conf = new Configuration();
    ShuffleVertexManagerBase manager = null;
    HashMap<String, EdgeProperty> mockInputVertices = new HashMap<String, EdgeProperty>();
    String r1 = "R1";
    EdgeProperty eProp1 = EdgeProperty.create(EdgeProperty.DataMovementType.SCATTER_GATHER, EdgeProperty.DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, OutputDescriptor.create("out"), InputDescriptor.create("in"));
    final String mockManagedVertexId = "R2";
    mockInputVertices.put(r1, eProp1);
    final VertexManagerPluginContext mockContext = mock(VertexManagerPluginContext.class);
    when(mockContext.getInputVertexEdgeProperties()).thenReturn(mockInputVertices);
    when(mockContext.getVertexName()).thenReturn(mockManagedVertexId);
    when(mockContext.getVertexNumTasks(mockManagedVertexId)).thenReturn(0);
    // check initialization
    manager = createManager(conf, mockContext, 0.001f, 0.001f);
    final List<Integer> scheduledTasks = Lists.newLinkedList();
    doAnswer(new ScheduledTasksAnswer(scheduledTasks)).when(mockContext).scheduleTasks(anyList());
    manager.onVertexStarted(emptyCompletions);
    manager.onVertexStateUpdated(new VertexStateUpdate(r1, VertexState.CONFIGURED));
    Assert.assertEquals(1, manager.bipartiteSources);
    Assert.assertEquals(0, manager.numBipartiteSourceTasksCompleted);
    Assert.assertEquals(0, manager.totalNumBipartiteSourceTasks);
    // no tasks scheduled
    Assert.assertEquals(0, manager.pendingTasks.size());
    Assert.assertEquals(0, scheduledTasks.size());
    verify(mockContext).doneReconfiguringVertex();
}

Example 55 with EdgeProperty

use of org.apache.tez.dag.api.EdgeProperty in project tez by apache.

the class TestShuffleVertexManagerBase method test_Tez1649_with_mixed_edges.

@Test(timeout = 5000)
public void test_Tez1649_with_mixed_edges() {
    Configuration conf = new Configuration();
    ShuffleVertexManagerBase manager = null;
    HashMap<String, EdgeProperty> mockInputVertices = new HashMap<String, EdgeProperty>();
    String r1 = "R1";
    EdgeProperty eProp1 = EdgeProperty.create(EdgeProperty.DataMovementType.SCATTER_GATHER, EdgeProperty.DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, OutputDescriptor.create("out"), InputDescriptor.create("in"));
    String m2 = "M2";
    EdgeProperty eProp2 = EdgeProperty.create(EdgeProperty.DataMovementType.BROADCAST, EdgeProperty.DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, OutputDescriptor.create("out"), InputDescriptor.create("in"));
    String m3 = "M3";
    EdgeProperty eProp3 = EdgeProperty.create(EdgeProperty.DataMovementType.BROADCAST, EdgeProperty.DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, OutputDescriptor.create("out"), InputDescriptor.create("in"));
    final String mockManagedVertexId = "R2";
    mockInputVertices.put(r1, eProp1);
    mockInputVertices.put(m2, eProp2);
    mockInputVertices.put(m3, eProp3);
    VertexManagerPluginContext mockContext = mock(VertexManagerPluginContext.class);
    when(mockContext.getInputVertexEdgeProperties()).thenReturn(mockInputVertices);
    when(mockContext.getVertexName()).thenReturn(mockManagedVertexId);
    when(mockContext.getVertexNumTasks(mockManagedVertexId)).thenReturn(3);
    when(mockContext.getVertexNumTasks(r1)).thenReturn(3);
    when(mockContext.getVertexNumTasks(m2)).thenReturn(3);
    when(mockContext.getVertexNumTasks(m3)).thenReturn(3);
    final List<Integer> scheduledTasks = Lists.newLinkedList();
    doAnswer(new ScheduledTasksAnswer(scheduledTasks)).when(mockContext).scheduleTasks(anyList());
    // check initialization
    manager = createManager(conf, mockContext, 0.001f, 0.001f);
    manager.onVertexStarted(emptyCompletions);
    Assert.assertTrue(manager.bipartiteSources == 1);
    manager.onVertexStateUpdated(new VertexStateUpdate(r1, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(m2, VertexState.CONFIGURED));
    // no tasks scheduled
    Assert.assertEquals(3, manager.pendingTasks.size());
    Assert.assertEquals(3, manager.totalNumBipartiteSourceTasks);
    Assert.assertEquals(0, manager.numBipartiteSourceTasksCompleted);
    // Send events for 2 tasks of r1.
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(r1, 0));
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(r1, 1));
    // no tasks scheduled
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    Assert.assertTrue(manager.totalNumBipartiteSourceTasks == 3);
    // Send an event for m2.
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(m2, 0));
    // no tasks scheduled
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    Assert.assertTrue(manager.totalNumBipartiteSourceTasks == 3);
    // Send an event for m3.
    manager.onVertexStateUpdated(new VertexStateUpdate(m3, VertexState.CONFIGURED));
    // all tasks scheduled
    Assert.assertTrue(manager.pendingTasks.size() == 0);
    Assert.assertTrue(scheduledTasks.size() == 3);
    // Scenario when numBipartiteSourceTasksCompleted == totalNumBipartiteSourceTasks.
    // Still, wait for a configuration to be completed from other edges
    scheduledTasks.clear();
    manager = createManager(conf, mockContext, 0.001f, 0.001f);
    manager.onVertexStarted(emptyCompletions);
    manager.onVertexStateUpdated(new VertexStateUpdate(r1, VertexState.CONFIGURED));
    when(mockContext.getInputVertexEdgeProperties()).thenReturn(mockInputVertices);
    when(mockContext.getVertexName()).thenReturn(mockManagedVertexId);
    when(mockContext.getVertexNumTasks(mockManagedVertexId)).thenReturn(3);
    when(mockContext.getVertexNumTasks(r1)).thenReturn(3);
    when(mockContext.getVertexNumTasks(m2)).thenReturn(3);
    when(mockContext.getVertexNumTasks(m3)).thenReturn(3);
    // no tasks scheduled
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    Assert.assertTrue(manager.totalNumBipartiteSourceTasks == 3);
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(r1, 0));
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(r1, 1));
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(r1, 2));
    // Tasks from non-scatter edges of m2 and m3 are not complete.
    // no tasks scheduled
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    manager.onVertexStateUpdated(new VertexStateUpdate(m2, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(m3, VertexState.CONFIGURED));
    // Got an event from other edges. Schedule all
    // all tasks scheduled
    Assert.assertTrue(manager.pendingTasks.size() == 0);
    Assert.assertTrue(scheduledTasks.size() == 3);
    // try with a zero task vertex (with non-scatter-gather edges)
    scheduledTasks.clear();
    manager = createManager(conf, mockContext, 0.001f, 0.001f);
    manager.onVertexStarted(emptyCompletions);
    when(mockContext.getInputVertexEdgeProperties()).thenReturn(mockInputVertices);
    when(mockContext.getVertexName()).thenReturn(mockManagedVertexId);
    when(mockContext.getVertexNumTasks(mockManagedVertexId)).thenReturn(3);
    // scatter gather
    when(mockContext.getVertexNumTasks(r1)).thenReturn(3);
    // broadcast
    when(mockContext.getVertexNumTasks(m2)).thenReturn(0);
    // broadcast
    when(mockContext.getVertexNumTasks(m3)).thenReturn(3);
    manager = createManager(conf, mockContext, 0.001f, 0.001f);
    manager.onVertexStarted(emptyCompletions);
    manager.onVertexStateUpdated(new VertexStateUpdate(r1, VertexState.CONFIGURED));
    // no tasks scheduled
    Assert.assertEquals(3, manager.pendingTasks.size());
    Assert.assertEquals(3, manager.totalNumBipartiteSourceTasks);
    Assert.assertEquals(0, manager.numBipartiteSourceTasksCompleted);
    // Send 2 events for tasks of r1.
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(r1, 0));
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(r1, 1));
    // no tasks scheduled
    Assert.assertTrue(manager.pendingTasks.size() == 3);
    Assert.assertTrue(scheduledTasks.size() == 0);
    // event from m3 triggers scheduling
    manager.onVertexStateUpdated(new VertexStateUpdate(m3, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(m2, VertexState.CONFIGURED));
    // all tasks scheduled
    Assert.assertTrue(manager.pendingTasks.size() == 0);
    Assert.assertTrue(scheduledTasks.size() == 3);
    // try with all zero task vertices in non-SG edges
    scheduledTasks.clear();
    manager = createManager(conf, mockContext, 0.001f, 0.001f);
    manager.onVertexStarted(emptyCompletions);
    when(mockContext.getInputVertexEdgeProperties()).thenReturn(mockInputVertices);
    when(mockContext.getVertexName()).thenReturn(mockManagedVertexId);
    when(mockContext.getVertexNumTasks(mockManagedVertexId)).thenReturn(3);
    // scatter gather
    when(mockContext.getVertexNumTasks(r1)).thenReturn(3);
    // broadcast
    when(mockContext.getVertexNumTasks(m2)).thenReturn(0);
    // broadcast
    when(mockContext.getVertexNumTasks(m3)).thenReturn(0);
    // Send 1 events for tasks of r1.
    manager.onVertexStateUpdated(new VertexStateUpdate(r1, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(m3, VertexState.CONFIGURED));
    manager.onVertexStateUpdated(new VertexStateUpdate(m2, VertexState.CONFIGURED));
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(r1, 0));
    // all tasks scheduled
    Assert.assertTrue(manager.pendingTasks.size() == 0);
    Assert.assertTrue(scheduledTasks.size() == 3);
}