Package org.apache.tez.dag.api

Examples of org.apache.tez.dag.api.EdgeManagerPlugin


  @Test(timeout = 5000)
  public void testSetCustomEdgeManager() throws UnsupportedEncodingException {
    initAllVertices(VertexState.INITED);
    Edge edge = edges.get("e4");
    EdgeManagerPlugin em = edge.getEdgeManager();
    EdgeManagerForTest originalEm = (EdgeManagerForTest) em;
    Assert.assertTrue(Arrays.equals(edgePayload, originalEm.getEdgeManagerContext()
        .getUserPayload().deepCopyAsArray()));

    UserPayload userPayload = UserPayload.create(ByteBuffer.wrap(new String("foo").getBytes()));
    EdgeManagerPluginDescriptor edgeManagerDescriptor =
        EdgeManagerPluginDescriptor.create(EdgeManagerForTest.class.getName());
    edgeManagerDescriptor.setUserPayload(userPayload);

    Vertex v3 = vertices.get("vertex3");
    Vertex v5 = vertices.get("vertex5"); // Vertex5 linked to v3 (v3 src, v5
                                         // dest)

    Map<String, EdgeManagerPluginDescriptor> edgeManagerDescriptors =
        Collections.singletonMap(v3.getName(), edgeManagerDescriptor);
    Assert.assertTrue(v5.setParallelism(v5.getTotalTasks() - 1, null,
        edgeManagerDescriptors, null)); // Must decrease.

    VertexImpl v5Impl = (VertexImpl) v5;

    EdgeManagerPlugin modifiedEdgeManager = v5Impl.sourceVertices.get(v3)
        .getEdgeManager();
    Assert.assertNotNull(modifiedEdgeManager);
    Assert.assertTrue(modifiedEdgeManager instanceof EdgeManagerForTest);

    // Ensure initialize() is called with the correct payload
View Full Code Here


              @Override
              public int getDestinationVertexNumTasks() {
                return 0;
              }
            };
            EdgeManagerPlugin edgeManager = ReflectionUtils
                .createClazzInstance(entry.getValue().getClassName(),
                    new Class[]{EdgeManagerPluginContext.class}, new Object[]{emContext});
            edgeManager.initialize();
            newEdgeManagers.put(entry.getKey(), edgeManager);
          }
          return null;
      }}).when(mockContext).setVertexParallelism(eq(2), any(VertexLocationHint.class), anyMap(), anyMap());
   
    // source vertices have 0 tasks. immediate start of all managed tasks
    when(mockContext.getVertexNumTasks(mockSrcVertexId1)).thenReturn(0);
    when(mockContext.getVertexNumTasks(mockSrcVertexId2)).thenReturn(0);
    when(mockContext.getVertexNumTasks(mockSrcVertexId3)).thenReturn(1);

    manager.onVertexStarted(null);
    Assert.assertTrue(manager.pendingTasks.isEmpty());
    Assert.assertTrue(scheduledTasks.size() == 4); // all tasks scheduled
    scheduledTasks.clear();
   
    when(mockContext.getVertexNumTasks(mockSrcVertexId1)).thenReturn(2);
    when(mockContext.getVertexNumTasks(mockSrcVertexId2)).thenReturn(2);

    ByteBuffer payload =
        VertexManagerEventPayloadProto.newBuilder().setOutputSize(5000L).build().toByteString().asReadOnlyByteBuffer();
    VertexManagerEvent vmEvent = VertexManagerEvent.create("Vertex", payload);
    // parallelism not change due to large data size
    manager = createManager(conf, mockContext, 0.1f, 0.1f);
    manager.onVertexStarted(null);
    Assert.assertTrue(manager.pendingTasks.size() == 4); // no tasks scheduled
    Assert.assertTrue(manager.totalNumSourceTasks == 4);
    manager.onVertexManagerEventReceived(vmEvent);
    manager.onSourceTaskCompleted(mockSrcVertexId1, new Integer(0));
    // managedVertex tasks reduced
    verify(mockContext, times(0)).setVertexParallelism(anyInt(), any(VertexLocationHint.class), anyMap(), anyMap());
    Assert.assertEquals(0, manager.pendingTasks.size()); // all tasks scheduled
    Assert.assertEquals(4, scheduledTasks.size());
    Assert.assertEquals(1, manager.numSourceTasksCompleted);
    Assert.assertEquals(5000L, manager.completedSourceTasksOutputSize);
   
   
    // parallelism changed due to small data size
    scheduledTasks.clear();
    payload =
        VertexManagerEventPayloadProto.newBuilder().setOutputSize(500L).build().toByteString().asReadOnlyByteBuffer();
    vmEvent = VertexManagerEvent.create("Vertex", payload);
   
    manager = createManager(conf, mockContext, 0.5f, 0.5f);
    manager.onVertexStarted(null);
    Assert.assertEquals(4, manager.pendingTasks.size()); // no tasks scheduled
    Assert.assertEquals(4, manager.totalNumSourceTasks);
    // task completion from non-bipartite stage does nothing
    manager.onSourceTaskCompleted(mockSrcVertexId3, new Integer(0));
    Assert.assertEquals(4, manager.pendingTasks.size()); // no tasks scheduled
    Assert.assertEquals(4, manager.totalNumSourceTasks);
    Assert.assertEquals(0, manager.numSourceTasksCompleted);
    manager.onVertexManagerEventReceived(vmEvent);
    manager.onSourceTaskCompleted(mockSrcVertexId1, new Integer(0));
    Assert.assertEquals(4, manager.pendingTasks.size());
    Assert.assertEquals(0, scheduledTasks.size()); // no tasks scheduled
    Assert.assertEquals(1, manager.numSourceTasksCompleted);
    Assert.assertEquals(1, manager.numVertexManagerEventsReceived);
    Assert.assertEquals(500L, manager.completedSourceTasksOutputSize);
    // ignore duplicate completion
    manager.onSourceTaskCompleted(mockSrcVertexId1, new Integer(0));
    Assert.assertEquals(4, manager.pendingTasks.size());
    Assert.assertEquals(0, scheduledTasks.size()); // no tasks scheduled
    Assert.assertEquals(1, manager.numSourceTasksCompleted);
    Assert.assertEquals(500L, manager.completedSourceTasksOutputSize);
   
    manager.onVertexManagerEventReceived(vmEvent);
    manager.onSourceTaskCompleted(mockSrcVertexId1, new Integer(1));
    // managedVertex tasks reduced
    verify(mockContext).setVertexParallelism(eq(2), any(VertexLocationHint.class), anyMap(), anyMap());
    Assert.assertEquals(2, newEdgeManagers.size());
    // TODO improve tests for parallelism
    Assert.assertEquals(0, manager.pendingTasks.size()); // all tasks scheduled
    Assert.assertEquals(2, scheduledTasks.size());
    Assert.assertTrue(scheduledTasks.contains(new Integer(0)));
    Assert.assertTrue(scheduledTasks.contains(new Integer(1)));
    Assert.assertEquals(2, manager.numSourceTasksCompleted);
    Assert.assertEquals(2, manager.numVertexManagerEventsReceived);
    Assert.assertEquals(1000L, manager.completedSourceTasksOutputSize);
   
    // more completions dont cause recalculation of parallelism
    manager.onSourceTaskCompleted(mockSrcVertexId2, new Integer(0));
    verify(mockContext).setVertexParallelism(eq(2), any(VertexLocationHint.class), anyMap(), anyMap());
    Assert.assertEquals(2, newEdgeManagers.size());
   
    EdgeManagerPlugin edgeManager = newEdgeManagers.values().iterator().next();
    Map<Integer, List<Integer>> targets = Maps.newHashMap();
    DataMovementEvent dmEvent = DataMovementEvent.create(1, ByteBuffer.wrap(new byte[0]));
    // 4 source task outputs - same as original number of partitions
    Assert.assertEquals(4, edgeManager.getNumSourceTaskPhysicalOutputs(0));
    // 4 destination task inputs - 2 source tasks + 2 merged partitions
    Assert.assertEquals(4, edgeManager.getNumDestinationTaskPhysicalInputs(0));
    edgeManager.routeDataMovementEventToDestination(dmEvent, 1, dmEvent.getSourceIndex(), targets);
    Assert.assertEquals(1, targets.size());
    Map.Entry<Integer, List<Integer>> e = targets.entrySet().iterator().next();
    Assert.assertEquals(0, e.getKey().intValue());
    Assert.assertEquals(1, e.getValue().size());
    Assert.assertEquals(3, e.getValue().get(0).intValue());
    targets.clear();
    dmEvent = DataMovementEvent.create(2, ByteBuffer.wrap(new byte[0]));
    edgeManager.routeDataMovementEventToDestination(dmEvent, 0, dmEvent.getSourceIndex(), targets);
    Assert.assertEquals(1, targets.size());
    e = targets.entrySet().iterator().next();
    Assert.assertEquals(1, e.getKey().intValue());
    Assert.assertEquals(1, e.getValue().size());
    Assert.assertEquals(0, e.getValue().get(0).intValue());
    targets.clear();
    edgeManager.routeInputSourceTaskFailedEventToDestination(2, targets);
    Assert.assertEquals(2, targets.size());
    for (Map.Entry<Integer, List<Integer>> entry : targets.entrySet()) {
      Assert.assertTrue(entry.getKey().intValue() == 0 || entry.getKey().intValue() == 1);
      Assert.assertEquals(2, entry.getValue().size());
      Assert.assertEquals(4, entry.getValue().get(0).intValue());
View Full Code Here

TOP

Related Classes of org.apache.tez.dag.api.EdgeManagerPlugin

Copyright © 2018 www.massapicom. All rights reserved.
All source code are property of their respective owners. Java is a trademark of Sun Microsystems, Inc and owned by ORACLE Inc. Contact coftware#gmail.com.