Examples with TaskAttemptIdentifier org.apache.tez.runtime.api.TaskAttemptIdentifier used on opensource projects

Example 1 with TaskAttemptIdentifier

use of org.apache.tez.runtime.api.TaskAttemptIdentifier in project tez by apache.

the class TestFairCartesianProductVertexManager method testZeroSrcTask.

@Test(timeout = 5000)
public void testZeroSrcTask() throws Exception {
    ctx = mock(VertexManagerPluginContext.class);
    vertexManager = new FairCartesianProductVertexManager(ctx);
    when(ctx.getVertexNumTasks("v0")).thenReturn(2);
    when(ctx.getVertexNumTasks("v1")).thenReturn(0);
    CartesianProductConfigProto.Builder builder = CartesianProductConfigProto.newBuilder();
    builder.setIsPartitioned(false).addSources("v0").addSources("v1").addNumChunks(2).addNumChunks(3).setMaxParallelism(6);
    CartesianProductConfigProto config = builder.build();
    Map<String, EdgeProperty> edgePropertyMap = new HashMap<>();
    edgePropertyMap.put("v0", EdgeProperty.create(EdgeManagerPluginDescriptor.create(CartesianProductEdgeManager.class.getName()), null, null, null, null));
    edgePropertyMap.put("v1", EdgeProperty.create(EdgeManagerPluginDescriptor.create(CartesianProductEdgeManager.class.getName()), null, null, null, null));
    when(ctx.getInputVertexEdgeProperties()).thenReturn(edgePropertyMap);
    vertexManager.initialize(config);
    vertexManager.onVertexStateUpdated(new VertexStateUpdate("v0", VertexState.CONFIGURED));
    vertexManager.onVertexStateUpdated(new VertexStateUpdate("v1", VertexState.CONFIGURED));
    vertexManager.onVertexStarted(new ArrayList<TaskAttemptIdentifier>());
    vertexManager.onSourceTaskCompleted(getTaId("v0", 0));
    vertexManager.onSourceTaskCompleted(getTaId("v0", 1));
}

Example 2 with TaskAttemptIdentifier

use of org.apache.tez.runtime.api.TaskAttemptIdentifier 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());
}

Example 3 with TaskAttemptIdentifier

use of org.apache.tez.runtime.api.TaskAttemptIdentifier in project tez by apache.

the class ShuffleVertexManagerBase method onVertexStarted.

@Override
public synchronized void onVertexStarted(List<TaskAttemptIdentifier> completions) {
    // examine edges after vertex started because until then these may not have been defined
    Map<String, EdgeProperty> inputs = getContext().getInputVertexEdgeProperties();
    for (Map.Entry<String, EdgeProperty> entry : inputs.entrySet()) {
        srcVertexInfo.put(entry.getKey(), createSourceVertexInfo(entry.getValue(), getContext().getVertexNumTasks(getContext().getVertexName())));
        // TODO what if derived class has already called this
        // register for status update from all source vertices
        getContext().registerForVertexStateUpdates(entry.getKey(), EnumSet.of(VertexState.CONFIGURED));
        if (entry.getValue().getDataMovementType() == DataMovementType.SCATTER_GATHER) {
            bipartiteSources++;
        }
    }
    onVertexStartedCheck();
    for (VertexStateUpdate stateUpdate : pendingStateUpdates) {
        handleVertexStateUpdate(stateUpdate);
    }
    pendingStateUpdates.clear();
    // track the tasks in this vertex
    updatePendingTasks();
    for (VertexManagerEvent vmEvent : pendingVMEvents) {
        handleVertexManagerEvent(vmEvent);
    }
    pendingVMEvents.clear();
    LOG.info("OnVertexStarted vertex: {} with {} source tasks and {} pending" + " tasks", getContext().getVertexName(), totalNumBipartiteSourceTasks, totalTasksToSchedule);
    if (completions != null) {
        for (TaskAttemptIdentifier attempt : completions) {
            onSourceTaskCompleted(attempt);
        }
    }
    onVertexStartedDone.set(true);
    // for the special case when source has 0 tasks or min fraction == 0
    processPendingTasks(null);
}

Example 4 with TaskAttemptIdentifier

use of org.apache.tez.runtime.api.TaskAttemptIdentifier in project tez by apache.

the class TestRootInputVertexManager method createTaskAttemptIdentifier.

public static TaskAttemptIdentifier createTaskAttemptIdentifier(String vName, int tId) {
    VertexIdentifier mockVertex = mock(VertexIdentifier.class);
    when(mockVertex.getName()).thenReturn(vName);
    TaskIdentifier mockTask = mock(TaskIdentifier.class);
    when(mockTask.getIdentifier()).thenReturn(tId);
    when(mockTask.getVertexIdentifier()).thenReturn(mockVertex);
    TaskAttemptIdentifier mockAttempt = mock(TaskAttemptIdentifier.class);
    when(mockAttempt.getIdentifier()).thenReturn(0);
    when(mockAttempt.getTaskIdentifier()).thenReturn(mockTask);
    return mockAttempt;
}

Example 5 with TaskAttemptIdentifier

use of org.apache.tez.runtime.api.TaskAttemptIdentifier in project tez by apache.

the class FairCartesianProductVertexManager method onVertexStarted.

@Override
public synchronized void onVertexStarted(List<TaskAttemptIdentifier> completions) throws Exception {
    vertexStarted = true;
    if (completions != null) {
        LOG.info("OnVertexStarted with " + completions.size() + " completed source task");
        for (TaskAttemptIdentifier attempt : completions) {
            addCompletedSrcTaskToProcess(attempt);
        }
    }
    tryScheduleTasks();
}