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

Example 11 with EdgeProperty

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

the class TestShuffleVertexManagerUtils method createVertexManagerContext.

VertexManagerPluginContext createVertexManagerContext(String mockSrcVertexId1, int numTasksSrcVertexId1, String mockSrcVertexId2, int numTasksSrcVertexId2, String mockSrcVertexId3, int numTasksSrcVertexId3, String mockManagedVertexId, int numTasksmockManagedVertexId, List<Integer> scheduledTasks, Map<String, EdgeManagerPlugin> newEdgeManagers) {
    HashMap<String, EdgeProperty> mockInputVertices = new HashMap<String, EdgeProperty>();
    EdgeProperty eProp1 = EdgeProperty.create(EdgeProperty.DataMovementType.SCATTER_GATHER, EdgeProperty.DataSourceType.PERSISTED, EdgeProperty.SchedulingType.SEQUENTIAL, OutputDescriptor.create("out"), InputDescriptor.create("in"));
    EdgeProperty eProp2 = EdgeProperty.create(EdgeProperty.DataMovementType.SCATTER_GATHER, EdgeProperty.DataSourceType.PERSISTED, EdgeProperty.SchedulingType.SEQUENTIAL, OutputDescriptor.create("out"), InputDescriptor.create("in"));
    EdgeProperty eProp3 = EdgeProperty.create(EdgeProperty.DataMovementType.BROADCAST, EdgeProperty.DataSourceType.PERSISTED, EdgeProperty.SchedulingType.SEQUENTIAL, OutputDescriptor.create("out"), InputDescriptor.create("in"));
    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(mockSrcVertexId1)).thenReturn(numTasksSrcVertexId1);
    when(mockContext.getVertexNumTasks(mockSrcVertexId2)).thenReturn(numTasksSrcVertexId2);
    when(mockContext.getVertexNumTasks(mockSrcVertexId3)).thenReturn(numTasksSrcVertexId3);
    when(mockContext.getVertexNumTasks(mockManagedVertexId)).thenReturn(numTasksmockManagedVertexId);
    doAnswer(new ScheduledTasksAnswer(scheduledTasks)).when(mockContext).scheduleTasks(anyList());
    doAnswer(new reconfigVertexAnswer(mockContext, mockManagedVertexId, newEdgeManagers)).when(mockContext).reconfigureVertex(anyInt(), any(VertexLocationHint.class), anyMap());
    return mockContext;
}

Example 12 with EdgeProperty

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

the class TestVertexImpl method testVertexSetParallelismIncrease.

@Test(timeout = 5000)
public void testVertexSetParallelismIncrease() throws Exception {
    VertexImpl v3 = vertices.get("vertex3");
    v3.vertexReconfigurationPlanned();
    initAllVertices(VertexState.INITED);
    Assert.assertEquals(2, v3.getTotalTasks());
    Assert.assertEquals(2, v3.getTasks().size());
    VertexImpl v1 = vertices.get("vertex1");
    startVertex(vertices.get("vertex2"));
    startVertex(v1);
    EdgeManagerPluginDescriptor mockEdgeManagerDescriptor = EdgeManagerPluginDescriptor.create(EdgeManagerForTest.class.getName());
    EdgeProperty edgeProp = EdgeProperty.create(mockEdgeManagerDescriptor, DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, OutputDescriptor.create("Out"), InputDescriptor.create("In"));
    Map<String, EdgeProperty> edgeManagerDescriptors = Collections.singletonMap(v1.getName(), edgeProp);
    v3.reconfigureVertex(10, null, edgeManagerDescriptors);
    v3.doneReconfiguringVertex();
    assertTrue(v3.sourceVertices.get(v1).getEdgeManager() instanceof EdgeManagerForTest);
    checkTasks(v3, 10);
}

Example 13 with EdgeProperty

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

the class TestVertexImpl method setupPostDagCreation.

@SuppressWarnings({ "unchecked", "rawtypes" })
public void setupPostDagCreation() throws TezException {
    String dagName = "dag0";
    // dispatcher may be created multiple times (setupPostDagCreation may be called multiples)
    if (dispatcher != null) {
        dispatcher.stop();
    }
    dispatcher = new DrainDispatcher();
    appContext = mock(AppContext.class);
    when(appContext.getHadoopShim()).thenReturn(new DefaultHadoopShim());
    when(appContext.getContainerLauncherName(anyInt())).thenReturn(TezConstants.getTezYarnServicePluginName());
    thh = mock(TaskHeartbeatHandler.class);
    historyEventHandler = mock(HistoryEventHandler.class);
    TaskSchedulerManager taskScheduler = mock(TaskSchedulerManager.class);
    UserGroupInformation ugi;
    try {
        ugi = UserGroupInformation.getCurrentUser();
    } catch (IOException e) {
        throw new RuntimeException(e);
    }
    DAG dag = mock(DAG.class);
    doReturn(ugi).when(dag).getDagUGI();
    doReturn(dagName).when(dag).getName();
    Map<String, LocalResource> localResources = new HashMap<>();
    for (PlanLocalResource planLR : dagPlan.getLocalResourceList()) {
        localResources.put(planLR.getName(), DagTypeConverters.convertPlanLocalResourceToLocalResource(planLR));
    }
    when(dag.getLocalResources()).thenReturn(localResources);
    doReturn(appAttemptId).when(appContext).getApplicationAttemptId();
    doReturn(appAttemptId.getApplicationId()).when(appContext).getApplicationID();
    doReturn(dag).when(appContext).getCurrentDAG();
    execService = mock(ListeningExecutorService.class);
    final ListenableFuture<Void> mockFuture = mock(ListenableFuture.class);
    Mockito.doAnswer(new Answer() {

        public ListenableFuture<Void> answer(InvocationOnMock invocation) {
            Object[] args = invocation.getArguments();
            CallableEvent e = (CallableEvent) args[0];
            dispatcher.getEventHandler().handle(e);
            return mockFuture;
        }
    }).when(execService).submit((Callable<Void>) any());
    MockClock clock = new MockClock();
    doReturn(execService).when(appContext).getExecService();
    doReturn(conf).when(appContext).getAMConf();
    doReturn(new Credentials()).when(dag).getCredentials();
    doReturn(DAGPlan.getDefaultInstance()).when(dag).getJobPlan();
    doReturn(dagId).when(appContext).getCurrentDAGID();
    doReturn(dagId).when(dag).getID();
    doReturn(taskScheduler).when(appContext).getTaskScheduler();
    doReturn(Resource.newInstance(102400, 60)).when(taskScheduler).getTotalResources(0);
    doReturn(historyEventHandler).when(appContext).getHistoryHandler();
    doReturn(dispatcher.getEventHandler()).when(appContext).getEventHandler();
    doReturn(clock).when(appContext).getClock();
    vertexGroups = Maps.newHashMap();
    for (PlanVertexGroupInfo groupInfo : dagPlan.getVertexGroupsList()) {
        vertexGroups.put(groupInfo.getGroupName(), new VertexGroupInfo(groupInfo));
    }
    // updateTracker may be created multiple times (setupPostDagCreation may be called multiples)
    if (updateTracker != null) {
        updateTracker.stop();
    }
    updateTracker = new StateChangeNotifierForTest(appContext.getCurrentDAG());
    setupVertices();
    when(dag.getVertex(any(TezVertexID.class))).thenAnswer(new Answer<Vertex>() {

        @Override
        public Vertex answer(InvocationOnMock invocation) throws Throwable {
            Object[] args = invocation.getArguments();
            if (args.length != 1) {
                return null;
            }
            TezVertexID vId = (TezVertexID) args[0];
            return vertexIdMap.get(vId);
        }
    });
    when(dag.getVertex(any(String.class))).thenAnswer(new Answer<Vertex>() {

        @Override
        public Vertex answer(InvocationOnMock invocation) throws Throwable {
            Object[] args = invocation.getArguments();
            if (args.length != 1) {
                return null;
            }
            String vId = (String) args[0];
            return vertices.get(vId);
        }
    });
    // TODO - this test logic is tightly linked to impl DAGImpl code.
    edges = new HashMap<String, Edge>();
    for (EdgePlan edgePlan : dagPlan.getEdgeList()) {
        EdgeProperty edgeProperty = DagTypeConverters.createEdgePropertyMapFromDAGPlan(edgePlan);
        edges.put(edgePlan.getId(), new Edge(edgeProperty, dispatcher.getEventHandler(), conf));
    }
    parseVertexEdges();
    for (Edge edge : edges.values()) {
        edge.initialize();
    }
    dispatcher.register(CallableEventType.class, new CallableEventDispatcher());
    taskAttemptEventDispatcher = new TaskAttemptEventDispatcher();
    dispatcher.register(TaskAttemptEventType.class, taskAttemptEventDispatcher);
    taskEventDispatcher = new TaskEventDispatcher();
    dispatcher.register(TaskEventType.class, taskEventDispatcher);
    vertexEventDispatcher = new VertexEventDispatcher();
    dispatcher.register(VertexEventType.class, vertexEventDispatcher);
    dagEventDispatcher = new DagEventDispatcher();
    dispatcher.register(DAGEventType.class, dagEventDispatcher);
    amSchedulerEventDispatcher = new AMSchedulerEventDispatcher();
    dispatcher.register(AMSchedulerEventType.class, amSchedulerEventDispatcher);
    dispatcher.init(conf);
    dispatcher.start();
}

Example 14 with EdgeProperty

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

the class TestVertexImpl method testStartWithUninitializedCustomEdge.

@Test(timeout = 5000)
public void testStartWithUninitializedCustomEdge() throws Exception {
    // Race when a source vertex manages to start before the target vertex has
    // been initialized
    setupPreDagCreation();
    dagPlan = createSamplerDAGPlan(true);
    setupPostDagCreation();
    VertexImpl vA = vertices.get("A");
    VertexImpl vB = vertices.get("B");
    VertexImpl vC = vertices.get("C");
    dispatcher.getEventHandler().handle(new VertexEvent(vA.getVertexId(), VertexEventType.V_INIT));
    dispatcher.getEventHandler().handle(new VertexEvent(vA.getVertexId(), VertexEventType.V_START));
    dispatcher.await();
    Assert.assertEquals(VertexState.INITIALIZING, vA.getState());
    Assert.assertEquals(VertexState.INITIALIZING, vB.getState());
    Assert.assertEquals(VertexState.INITIALIZING, vC.getState());
    // setting the edge manager should vA to start
    EdgeManagerPluginDescriptor mockEdgeManagerDescriptor = EdgeManagerPluginDescriptor.create(EdgeManagerForTest.class.getName());
    Edge e = vC.sourceVertices.get(vA);
    Assert.assertNull(e.getEdgeManager());
    e.setCustomEdgeManager(mockEdgeManagerDescriptor);
    dispatcher.await();
    Assert.assertEquals(VertexState.RUNNING, vA.getState());
    Assert.assertEquals(VertexState.INITIALIZING, vB.getState());
    Assert.assertEquals(VertexState.INITIALIZING, vC.getState());
    EdgeProperty edgeProp = EdgeProperty.create(mockEdgeManagerDescriptor, DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, OutputDescriptor.create("Out"), InputDescriptor.create("In"));
    Map<String, EdgeProperty> edges = Maps.newHashMap();
    edges.put("B", edgeProp);
    vC.reconfigureVertex(2, vertexLocationHint, edges);
    dispatcher.await();
    Assert.assertEquals(VertexState.RUNNING, vA.getState());
    Assert.assertEquals(VertexState.RUNNING, vB.getState());
    Assert.assertEquals(VertexState.RUNNING, vC.getState());
    Assert.assertNotNull(vA.getTask(0));
    Assert.assertNotNull(vB.getTask(0));
    Assert.assertNotNull(vC.getTask(0));
}

Example 15 with EdgeProperty

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