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

Example 1 with ScheduleTaskRequest

use of org.apache.tez.dag.api.VertexManagerPluginContext.ScheduleTaskRequest in project tez by apache.

the class TestCartesianProductVertexManagerPartitioned method testScheduling.

@Test(timeout = 5000)
public void testScheduling() throws Exception {
    vertexManager.onVertexStarted(null);
    vertexManager.onVertexStateUpdated(new VertexStateUpdate("v0", VertexState.CONFIGURED));
    vertexManager.onVertexStateUpdated(new VertexStateUpdate("v1", VertexState.CONFIGURED));
    vertexManager.onSourceTaskCompleted(allCompletions.get(0));
    vertexManager.onSourceTaskCompleted(allCompletions.get(1));
    verify(context, never()).scheduleTasks(Matchers.<List<ScheduleTaskRequest>>any());
    List<ScheduleTaskRequest> scheduleTaskRequests;
    vertexManager.onSourceTaskCompleted(allCompletions.get(2));
    // shouldn't start schedule because broadcast src is not in RUNNING state
    verify(context, never()).scheduleTasks(Matchers.<List<ScheduleTaskRequest>>any());
    vertexManager.onVertexStateUpdated(new VertexStateUpdate("v2", VertexState.RUNNING));
    verify(context, times(1)).scheduleTasks(scheduleTaskRequestCaptor.capture());
    scheduleTaskRequests = scheduleTaskRequestCaptor.getValue();
    assertEquals(1, scheduleTaskRequests.size());
    assertEquals(0, scheduleTaskRequests.get(0).getTaskIndex());
    // completion from broadcast src shouldn't matter
    vertexManager.onSourceTaskCompleted(allCompletions.get(8));
    verify(context, times(1)).scheduleTasks(scheduleTaskRequestCaptor.capture());
    for (int i = 3; i < 6; i++) {
        vertexManager.onSourceTaskCompleted(allCompletions.get(i));
        verify(context, times(i - 1)).scheduleTasks(scheduleTaskRequestCaptor.capture());
        scheduleTaskRequests = scheduleTaskRequestCaptor.getValue();
        assertEquals(1, scheduleTaskRequests.size());
        assertEquals(i - 2, scheduleTaskRequests.get(0).getTaskIndex());
    }
    for (int i = 6; i < 8; i++) {
        vertexManager.onSourceTaskCompleted(allCompletions.get(i));
        verify(context, times(4)).scheduleTasks(Matchers.<List<ScheduleTaskRequest>>any());
    }
}

Example 2 with ScheduleTaskRequest

use of org.apache.tez.dag.api.VertexManagerPluginContext.ScheduleTaskRequest in project tez by apache.

the class TestVertexImpl method testVertexGetTAAttemptsObsoletionWithPendingRoutes.

@Test(timeout = 5000)
public void testVertexGetTAAttemptsObsoletionWithPendingRoutes() throws Exception {
    initAllVertices(VertexState.INITED);
    VertexImpl v1 = vertices.get("vertex1");
    startVertex(v1);
    VertexImpl v2 = vertices.get("vertex2");
    startVertex(v2);
    VertexImpl v3 = vertices.get("vertex3");
    VertexImpl v4 = vertices.get("vertex4");
    List<ScheduleTaskRequest> taskList = new LinkedList<VertexManagerPluginContext.ScheduleTaskRequest>();
    // scheduling start to trigger edge routing to begin
    for (int i = 0; i < v4.getTotalTasks(); ++i) {
        taskList.add(ScheduleTaskRequest.create(i, null));
    }
    v4.scheduleTasks(taskList);
    Assert.assertEquals(VertexState.RUNNING, v4.getState());
    Assert.assertEquals(1, v4.sourceVertices.size());
    Edge e = v4.sourceVertices.get(v3);
    TezTaskAttemptID v3TaId = TezTaskAttemptID.getInstance(TezTaskID.getInstance(v3.getVertexId(), 0), 0);
    TezTaskAttemptID v4TaId = TezTaskAttemptID.getInstance(TezTaskID.getInstance(v4.getVertexId(), 0), 0);
    for (int i = 0; i < 11; ++i) {
        v4.handle(new VertexEventRouteEvent(v4.getVertexId(), Collections.singletonList(new TezEvent(DataMovementEvent.create(0, null), new EventMetaData(EventProducerConsumerType.OUTPUT, v3.getName(), v3.getName(), v3TaId)))));
    }
    dispatcher.await();
    // verify all events have been are in taskEvents
    Assert.assertEquals(11, v4.getOnDemandRouteEvents().size());
    TaskAttemptEventInfo eventInfo;
    EdgeManagerPluginOnDemand mockPlugin = mock(EdgeManagerPluginOnDemand.class);
    EventRouteMetadata mockFailedRoute = EventRouteMetadata.create(1, new int[] { 0 });
    e.edgeManager = mockPlugin;
    when(mockPlugin.routeInputSourceTaskFailedEventToDestination(anyInt(), anyInt())).thenReturn(mockFailedRoute);
    // return more events that dont evenly fit in max size
    EventRouteMetadata mockRoute = EventRouteMetadata.create(2, new int[] { 0, 0 });
    when(mockPlugin.routeDataMovementEventToDestination(anyInt(), anyInt(), anyInt())).thenReturn(mockRoute);
    int fromEventId = 0;
    eventInfo = v4.getTaskAttemptTezEvents(v4TaId, fromEventId, 0, 5);
    fromEventId = eventInfo.getNextFromEventId();
    // 0-1 events expanded and fit, 2nd event has pending routes
    Assert.assertEquals(2, fromEventId);
    Assert.assertEquals(5, eventInfo.getEvents().size());
    // send an input failed event
    v4.handle(new VertexEventRouteEvent(v4.getVertexId(), Collections.singletonList(new TezEvent(InputFailedEvent.create(0, 0), new EventMetaData(EventProducerConsumerType.OUTPUT, v3.getName(), v3.getName(), v3TaId)))));
    // get only input failed event. all DM events obsoleted
    eventInfo = v4.getTaskAttemptTezEvents(v4TaId, fromEventId, 0, 5);
    fromEventId = eventInfo.getNextFromEventId();
    Assert.assertEquals(12, fromEventId);
    Assert.assertEquals(1, eventInfo.getEvents().size());
    Assert.assertEquals(EventType.INPUT_FAILED_EVENT, eventInfo.getEvents().get(0).getEventType());
}

Example 3 with ScheduleTaskRequest

use of org.apache.tez.dag.api.VertexManagerPluginContext.ScheduleTaskRequest in project tez by apache.

the class TestVertexImpl method testVertexWithInitializerSuccess.

@Test(timeout = 500000)
public void testVertexWithInitializerSuccess() throws Exception {
    useCustomInitializer = true;
    setupPreDagCreation();
    dagPlan = createDAGPlanWithInputInitializer("TestInputInitializer");
    setupPostDagCreation();
    VertexImplWithControlledInitializerManager v1 = (VertexImplWithControlledInitializerManager) vertices.get("vertex1");
    dispatcher.getEventHandler().handle(new VertexEvent(v1.getVertexId(), VertexEventType.V_INIT));
    dispatcher.await();
    Assert.assertEquals(VertexState.INITIALIZING, v1.getState());
    RootInputInitializerManagerControlled initializerManager1 = v1.getRootInputInitializerManager();
    List<TaskLocationHint> v1Hints = createTaskLocationHints(5);
    initializerManager1.completeInputInitialization(0, 5, v1Hints);
    dispatcher.await();
    Assert.assertEquals(VertexState.INITED, v1.getState());
    Assert.assertEquals(5, v1.getTotalTasks());
    Assert.assertEquals(RootInputVertexManager.class.getName(), v1.getVertexManager().getPlugin().getClass().getName());
    for (int i = 0; i < v1Hints.size(); ++i) {
        Assert.assertEquals(v1Hints.get(i), v1.getTaskLocationHints()[i]);
    }
    Assert.assertEquals(true, initializerManager1.hasShutDown);
    for (int i = 0; i < 5; i++) {
        List<InputSpec> inputSpecs = v1.getInputSpecList(i);
        Assert.assertEquals(1, inputSpecs.size());
        Assert.assertEquals(1, inputSpecs.get(0).getPhysicalEdgeCount());
    }
    List<ScheduleTaskRequest> taskList = new LinkedList<VertexManagerPluginContext.ScheduleTaskRequest>();
    // scheduling start to trigger edge routing to begin
    for (int i = 0; i < v1.getTotalTasks(); ++i) {
        taskList.add(ScheduleTaskRequest.create(i, null));
    }
    v1.scheduleTasks(taskList);
    dispatcher.await();
    // check all tasks get their events
    for (int i = 0; i < v1.getTotalTasks(); ++i) {
        Assert.assertEquals(1, v1.getTaskAttemptTezEvents(TezTaskAttemptID.getInstance(v1.getTask(i).getTaskId(), 0), 0, 0, 100).getEvents().size());
    }
    VertexImplWithControlledInitializerManager v2 = (VertexImplWithControlledInitializerManager) vertices.get("vertex2");
    Assert.assertEquals(VertexState.INITIALIZING, v2.getState());
    // non-task events don't get buffered
    List<TezEvent> events = Lists.newLinkedList();
    TezTaskID t0_v1 = TezTaskID.getInstance(v1.getVertexId(), 0);
    TezTaskAttemptID ta0_t0_v1 = TezTaskAttemptID.getInstance(t0_v1, 0);
    events.add(new TezEvent(VertexManagerEvent.create("vertex2", ByteBuffer.wrap(new byte[0])), new EventMetaData(EventProducerConsumerType.PROCESSOR, "vertex1", "vertex2", ta0_t0_v1)));
    events.add(new TezEvent(InputDataInformationEvent.createWithSerializedPayload(0, ByteBuffer.wrap(new byte[0])), new EventMetaData(EventProducerConsumerType.INPUT, "vertex2", "NULL_VERTEX", null)));
    dispatcher.getEventHandler().handle(new VertexEventRouteEvent(v2.getVertexId(), events));
    dispatcher.await();
    RootInputInitializerManagerControlled initializerManager2 = v2.getRootInputInitializerManager();
    List<TaskLocationHint> v2Hints = createTaskLocationHints(10);
    initializerManager2.completeInputInitialization(0, 10, v2Hints);
    dispatcher.await();
    Assert.assertEquals(VertexState.INITED, v2.getState());
    Assert.assertEquals(10, v2.getTotalTasks());
    Assert.assertEquals(RootInputVertexManager.class.getName(), v2.getVertexManager().getPlugin().getClass().getName());
    for (int i = 0; i < v2Hints.size(); ++i) {
        Assert.assertEquals(v2Hints.get(i), v2.getTaskLocationHints()[i]);
    }
    Assert.assertEquals(true, initializerManager2.hasShutDown);
    // scheduling start to trigger edge routing to begin
    taskList = new LinkedList<VertexManagerPluginContext.ScheduleTaskRequest>();
    // scheduling start to trigger edge routing to begin
    for (int i = 0; i < v2.getTotalTasks(); ++i) {
        taskList.add(ScheduleTaskRequest.create(i, null));
    }
    v2.scheduleTasks(taskList);
    dispatcher.await();
    // check all tasks get their events
    for (int i = 0; i < v2.getTotalTasks(); ++i) {
        Assert.assertEquals(((i == 0) ? 2 : 1), v2.getTaskAttemptTezEvents(TezTaskAttemptID.getInstance(v2.getTask(i).getTaskId(), 0), 0, 0, 100).getEvents().size());
    }
    for (int i = 0; i < 10; i++) {
        List<InputSpec> inputSpecs = v1.getInputSpecList(i);
        Assert.assertEquals(1, inputSpecs.size());
        Assert.assertEquals(1, inputSpecs.get(0).getPhysicalEdgeCount());
    }
}

Example 4 with ScheduleTaskRequest

use of org.apache.tez.dag.api.VertexManagerPluginContext.ScheduleTaskRequest in project tez by apache.

the class TestVertexImpl method testVertexGetTAAttempts.

@Test(timeout = 5000)
public void testVertexGetTAAttempts() throws Exception {
    initAllVertices(VertexState.INITED);
    VertexImpl v1 = vertices.get("vertex1");
    startVertex(v1);
    VertexImpl v2 = vertices.get("vertex2");
    startVertex(v2);
    VertexImpl v3 = vertices.get("vertex3");
    VertexImpl v4 = vertices.get("vertex4");
    Assert.assertEquals(VertexState.RUNNING, v4.getState());
    Assert.assertEquals(1, v4.sourceVertices.size());
    Edge e = v4.sourceVertices.get(v3);
    TezTaskAttemptID v3TaId = TezTaskAttemptID.getInstance(TezTaskID.getInstance(v3.getVertexId(), 0), 0);
    TezTaskAttemptID v4TaId = TezTaskAttemptID.getInstance(TezTaskID.getInstance(v4.getVertexId(), 0), 0);
    for (int i = 0; i < 5; ++i) {
        v4.handle(new VertexEventRouteEvent(v4.getVertexId(), Collections.singletonList(new TezEvent(DataMovementEvent.create(0, null), new EventMetaData(EventProducerConsumerType.OUTPUT, v3.getName(), v3.getName(), v3TaId)))));
    }
    dispatcher.await();
    // verify all events have been put in pending.
    // this is not necessary after legacy routing has been removed
    Assert.assertEquals(5, v4.pendingTaskEvents.size());
    List<ScheduleTaskRequest> taskList = new LinkedList<VertexManagerPluginContext.ScheduleTaskRequest>();
    // scheduling start to trigger edge routing to begin
    for (int i = 0; i < v4.getTotalTasks(); ++i) {
        taskList.add(ScheduleTaskRequest.create(i, null));
    }
    v4.scheduleTasks(taskList);
    dispatcher.await();
    // verify all events have been moved to taskEvents
    Assert.assertEquals(5, v4.getOnDemandRouteEvents().size());
    for (int i = 5; i < 11; ++i) {
        v4.handle(new VertexEventRouteEvent(v4.getVertexId(), Collections.singletonList(new TezEvent(DataMovementEvent.create(0, null), new EventMetaData(EventProducerConsumerType.OUTPUT, v3.getName(), v3.getName(), v3TaId)))));
    }
    dispatcher.await();
    // verify all events have been are in taskEvents
    Assert.assertEquals(11, v4.getOnDemandRouteEvents().size());
    TaskAttemptEventInfo eventInfo;
    EdgeManagerPluginOnDemand mockPlugin = mock(EdgeManagerPluginOnDemand.class);
    EventRouteMetadata mockRoute = EventRouteMetadata.create(1, new int[] { 0 });
    e.edgeManager = mockPlugin;
    // source task id will not match. all events will return null
    when(mockPlugin.routeDataMovementEventToDestination(1, 0, 0)).thenReturn(mockRoute);
    eventInfo = v4.getTaskAttemptTezEvents(v4TaId, 0, 0, 1);
    // all events traversed
    Assert.assertEquals(11, eventInfo.getNextFromEventId());
    // no events
    Assert.assertEquals(0, eventInfo.getEvents().size());
    int fromEventId = 0;
    // source task id will match. all events will be returned
    // max events is respected.
    when(mockPlugin.routeDataMovementEventToDestination(anyInt(), anyInt(), anyInt())).thenReturn(mockRoute);
    for (int i = 0; i < 11; ++i) {
        eventInfo = v4.getTaskAttemptTezEvents(v4TaId, fromEventId, 0, 1);
        fromEventId = eventInfo.getNextFromEventId();
        Assert.assertEquals((i + 1), fromEventId);
        Assert.assertEquals(1, eventInfo.getEvents().size());
    }
    eventInfo = v4.getTaskAttemptTezEvents(v4TaId, fromEventId, 0, 1);
    // all events traversed
    Assert.assertEquals(11, eventInfo.getNextFromEventId());
    // no events
    Assert.assertEquals(0, eventInfo.getEvents().size());
    // ask for events with sufficient buffer. get all events in a single shot.
    fromEventId = 0;
    eventInfo = v4.getTaskAttemptTezEvents(v4TaId, fromEventId, 0, 100);
    fromEventId = eventInfo.getNextFromEventId();
    Assert.assertEquals(11, fromEventId);
    Assert.assertEquals(11, eventInfo.getEvents().size());
    // change max events to larger value. max events does not evenly divide total events
    fromEventId = 0;
    for (int i = 1; i <= 2; ++i) {
        eventInfo = v4.getTaskAttemptTezEvents(v4TaId, fromEventId, 0, 5);
        fromEventId = eventInfo.getNextFromEventId();
        Assert.assertEquals((i * 5), fromEventId);
        Assert.assertEquals(5, eventInfo.getEvents().size());
    }
    eventInfo = v4.getTaskAttemptTezEvents(v4TaId, fromEventId, 0, 5);
    // all events traversed
    Assert.assertEquals(11, eventInfo.getNextFromEventId());
    // remainder events
    Assert.assertEquals(1, eventInfo.getEvents().size());
    // return more events that dont evenly fit in max size
    mockRoute = EventRouteMetadata.create(2, new int[] { 0, 0 });
    when(mockPlugin.routeDataMovementEventToDestination(anyInt(), anyInt(), anyInt())).thenReturn(mockRoute);
    fromEventId = 0;
    int lastFromEventId = 0;
    for (int i = 1; i <= 4; ++i) {
        eventInfo = v4.getTaskAttemptTezEvents(v4TaId, fromEventId, 0, 5);
        fromEventId = eventInfo.getNextFromEventId();
        Assert.assertEquals((i % 2 > 0 ? (lastFromEventId += 2) : (lastFromEventId += 3)), fromEventId);
        Assert.assertEquals(5, eventInfo.getEvents().size());
    }
    eventInfo = v4.getTaskAttemptTezEvents(v4TaId, fromEventId, 0, 5);
    // all events traversed
    Assert.assertEquals(11, eventInfo.getNextFromEventId());
    // remainder events
    Assert.assertEquals(2, eventInfo.getEvents().size());
}

Example 5 with ScheduleTaskRequest

use of org.apache.tez.dag.api.VertexManagerPluginContext.ScheduleTaskRequest in project tez by apache.

the class TestImmediateStartVertexManager method testBasic.

@SuppressWarnings({ "unchecked", "rawtypes" })
@Test(timeout = 5000)
public void testBasic() {
    HashMap<String, EdgeProperty> mockInputVertices = new HashMap<String, EdgeProperty>();
    final String mockSrcVertexId1 = "Vertex1";
    EdgeProperty eProp1 = EdgeProperty.create(EdgeProperty.DataMovementType.SCATTER_GATHER, EdgeProperty.DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, OutputDescriptor.create("out"), InputDescriptor.create("in"));
    final String mockSrcVertexId2 = "Vertex2";
    EdgeProperty eProp2 = EdgeProperty.create(mock(EdgeManagerPluginDescriptor.class), EdgeProperty.DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, OutputDescriptor.create("out"), InputDescriptor.create("in"));
    final String mockSrcVertexId3 = "Vertex3";
    EdgeProperty eProp3 = EdgeProperty.create(EdgeProperty.DataMovementType.BROADCAST, EdgeProperty.DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, OutputDescriptor.create("out"), InputDescriptor.create("in"));
    final String mockManagedVertexId = "Vertex4";
    mockInputVertices.put(mockSrcVertexId1, eProp1);
    mockInputVertices.put(mockSrcVertexId2, eProp2);
    mockInputVertices.put(mockSrcVertexId3, eProp3);
    final VertexManagerPluginContext mockContext = mock(VertexManagerPluginContext.class);
    when(mockContext.getInputVertexEdgeProperties()).thenReturn(mockInputVertices);
    when(mockContext.getVertexName()).thenReturn(mockManagedVertexId);
    when(mockContext.getVertexNumTasks(mockManagedVertexId)).thenReturn(4);
    when(mockContext.getVertexNumTasks(mockSrcVertexId1)).thenReturn(2);
    when(mockContext.getVertexNumTasks(mockSrcVertexId2)).thenReturn(2);
    when(mockContext.getVertexNumTasks(mockSrcVertexId3)).thenReturn(2);
    final HashSet<Integer> scheduledTasks = new HashSet<Integer>();
    doAnswer(new Answer() {

        public Object answer(InvocationOnMock invocation) {
            Object[] args = invocation.getArguments();
            scheduledTasks.clear();
            List<ScheduleTaskRequest> tasks = (List<ScheduleTaskRequest>) args[0];
            for (ScheduleTaskRequest task : tasks) {
                scheduledTasks.add(task.getTaskIndex());
            }
            return null;
        }
    }).when(mockContext).scheduleTasks(anyList());
    List<TaskAttemptIdentifier> emptyCompletions = null;
    ImmediateStartVertexManager manager = new ImmediateStartVertexManager(mockContext);
    manager.initialize();
    manager.onVertexStarted(emptyCompletions);
    verify(mockContext, times(0)).scheduleTasks(anyList());
    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)).scheduleTasks(anyList());
    Assert.assertEquals(4, scheduledTasks.size());
    // simulate race between onVertexStarted and notifications
    scheduledTasks.clear();
    final ImmediateStartVertexManager raceManager = new ImmediateStartVertexManager(mockContext);
    doAnswer(new Answer() {

        public Object answer(InvocationOnMock invocation) throws Exception {
            raceManager.onVertexStateUpdated(new VertexStateUpdate((String) invocation.getArguments()[0], VertexState.CONFIGURED));
            scheduledTasks.clear();
            return null;
        }
    }).when(mockContext).registerForVertexStateUpdates(anyString(), anySet());
    raceManager.initialize();
    raceManager.onVertexStarted(emptyCompletions);
    verify(mockContext, times(2)).scheduleTasks(anyList());
    Assert.assertEquals(4, scheduledTasks.size());
}