org.apache.tez.dag.api.EdgeProperty Java Examples

@Test(timeout = 5000)
public void testOtherEdgeType() throws Exception {
  // forbid other custom edge
  edgePropertyMap.put("v2", customEdge);
  config = new CartesianProductConfig(new int[]{2, 3}, new String[]{"v0", "v1"}, null);
  try {
    vertexManager.initialize();
    assertTrue(false);
  } catch (Exception ignored) {}

  // broadcast edge should be allowed and other non-custom edge shouldn't be allowed
  for (DataMovementType type : DataMovementType.values()) {
    if (type == CUSTOM) {
      continue;
    }
    edgePropertyMap.put("v2", EdgeProperty.create(type, null, null, null, null));
    try {
      vertexManager.initialize();
      assertTrue(type == BROADCAST);
    } catch (Exception e) {
      assertTrue(type != BROADCAST);
    }
  }
}
 

private GroupInputEdge newGroupInputEdge(TezOperator fromOp,
        TezOperator toOp, VertexGroup from, Vertex to) throws IOException {

    EdgeProperty edgeProperty = newEdge(fromOp, toOp);

    String groupInputClass = ConcatenatedMergedKeyValueInput.class.getName();

    // In case of SCATTER_GATHER and UnorderedKVInput it will still be
    // ConcatenatedMergedKeyValueInput
    if(edgeProperty.getDataMovementType().equals(DataMovementType.SCATTER_GATHER)
            && edgeProperty.getEdgeDestination().getClassName().equals(OrderedGroupedKVInput.class.getName())) {
        groupInputClass = OrderedGroupedMergedKVInput.class.getName();
    }

    return GroupInputEdge.create(from, to, edgeProperty,
            InputDescriptor.create(groupInputClass).setUserPayload(edgeProperty.getEdgeDestination().getUserPayload())
                .setHistoryText(edgeProperty.getEdgeDestination().getHistoryText()));
}
 

@Test (timeout=2000)
public void testHistoryText() {
  OrderedPartitionedKVEdgeConfig.Builder builder =
      OrderedPartitionedKVEdgeConfig.newBuilder("KEY", "VALUE", "PARTITIONER");
  Configuration fromConf = new Configuration(false);
  fromConf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_CONVERT_USER_PAYLOAD_TO_HISTORY_TEXT,
      true);
  builder.setFromConfiguration(fromConf);

  OrderedPartitionedKVEdgeConfig kvEdgeConfig = builder.build();
  checkHistoryText(kvEdgeConfig.getInputHistoryText());
  checkHistoryText(kvEdgeConfig.getOutputHistoryText());

  EdgeProperty defaultEdgeProperty = builder.build().createDefaultEdgeProperty();
  checkHistoryText(defaultEdgeProperty.getEdgeDestination().getHistoryText());
  checkHistoryText(defaultEdgeProperty.getEdgeSource().getHistoryText());

  EdgeManagerPluginDescriptor descriptor = mock(EdgeManagerPluginDescriptor.class);
  EdgeProperty edgeProperty = builder.build().createDefaultCustomEdgeProperty(descriptor);
  checkHistoryText(edgeProperty.getEdgeDestination().getHistoryText());
  checkHistoryText(edgeProperty.getEdgeSource().getHistoryText());
}
 

@Test (timeout=2000)
public void testHistoryText() {
  UnorderedPartitionedKVEdgeConfig.Builder builder =
      UnorderedPartitionedKVEdgeConfig.newBuilder("KEY", "VALUE", "PARTITIONER");
  Configuration fromConf = new Configuration(false);
  fromConf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_CONVERT_USER_PAYLOAD_TO_HISTORY_TEXT,
      true);
  builder.setFromConfiguration(fromConf);

  UnorderedPartitionedKVEdgeConfig kvEdgeConfig = builder.build();

  checkHistoryText(kvEdgeConfig.getInputHistoryText());
  checkHistoryText(kvEdgeConfig.getOutputHistoryText());

  EdgeProperty defaultEdgeProperty = builder.build().createDefaultEdgeProperty();
  checkHistoryText(defaultEdgeProperty.getEdgeDestination().getHistoryText());
  checkHistoryText(defaultEdgeProperty.getEdgeSource().getHistoryText());

  EdgeManagerPluginDescriptor descriptor = mock(EdgeManagerPluginDescriptor.class);
  EdgeProperty edgeProperty = builder.build().createDefaultCustomEdgeProperty(descriptor);
  checkHistoryText(edgeProperty.getEdgeDestination().getHistoryText());
  checkHistoryText(edgeProperty.getEdgeSource().getHistoryText());

}
 

public synchronized void setCustomEdgeManager(EdgeManagerDescriptor descriptor) {
  EdgeProperty modifiedEdgeProperty =
      new EdgeProperty(descriptor,
          edgeProperty.getDataSourceType(),
          edgeProperty.getSchedulingType(),
          edgeProperty.getEdgeSource(),
          edgeProperty.getEdgeDestination());
  this.edgeProperty = modifiedEdgeProperty;
  boolean wasUnInitialized = (edgeManager == null);
  createEdgeManager();
  initialize();
  if (wasUnInitialized) {
    sendEvent(new VertexEventNullEdgeInitialized(sourceVertex.getVertexId(), this, destinationVertex));
    sendEvent(new VertexEventNullEdgeInitialized(destinationVertex.getVertexId(), this, sourceVertex));
  }
}
 

@Test(timeout = 5000)
public void testInvalidSourceTaskIndex() throws Exception {
  EventHandler mockEventHandler = mock(EventHandler.class);
  Edge edge = new Edge(EdgeProperty.create(
      EdgeManagerPluginDescriptor.create(CustomEdgeManagerWithInvalidReturnValue.class.getName())
        .setUserPayload(new CustomEdgeManagerWithInvalidReturnValue.EdgeManagerConfig(1,1,1,-1).toUserPayload()),
      DataSourceType.PERSISTED,
      SchedulingType.SEQUENTIAL,
      OutputDescriptor.create(""),
      InputDescriptor.create("")), mockEventHandler, new TezConfiguration());
  TezVertexID v1Id = createVertexID(1);
  TezVertexID v2Id = createVertexID(2);
  edge.setSourceVertex(mockVertex("v1", v1Id, new LinkedHashMap<TezTaskID, Task>()));
  edge.setDestinationVertex(mockVertex("v2", v2Id, new LinkedHashMap<TezTaskID, Task>()));
  edge.initialize();
  try {
    TezEvent ireEvent = new TezEvent(InputReadErrorEvent.create("diag", 0, 1),
        new EventMetaData(EventProducerConsumerType.INPUT, "v2", "v1",
            TezTaskAttemptID.getInstance(TezTaskID.getInstance(v2Id, 1), 1)));
    edge.sendTezEventToSourceTasks(ireEvent);
    Assert.fail();
  } catch (AMUserCodeException e) {
    e.printStackTrace();
    assertTrue(e.getCause().getMessage().contains("SourceTaskIndex should not be negative"));
  }
}
 

@Test(timeout = 5000)
public void testParallelismTwoSkewedSource() throws Exception {
  setupDAGVertexOnly(100, 10000, 10, 10);
  vertexManager.onVertexStateUpdated(new VertexStateUpdate("v0", VertexState.CONFIGURED));
  vertexManager.onVertexStateUpdated(new VertexStateUpdate("v1", VertexState.CONFIGURED));

  vertexManager.onVertexManagerEventReceived(getVMEvent(15000, "v0", 0));

  for (int i = 0; i < 30; i++) {
    vertexManager.onVertexManagerEventReceived(getVMEvent(1, "v1", i));
  }

  verify(ctx, times(1)).reconfigureVertex(
    eq(99), any(VertexLocationHint.class), edgePropertiesCaptor.capture());
  Map<String, EdgeProperty> edgeProperties = edgePropertiesCaptor.getValue();
  verifyEdgeProperties(edgeProperties.get("v0"), new String[]{"v0", "v1"},
    new int[]{99, 1}, 100);
}
 

@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));
}
 

public static DAG createDAG(String name, 
    Configuration conf) throws Exception {
  UserPayload payload = UserPayload.create(null);
  int taskCount = TEZ_SIMPLE_DAG_NUM_TASKS_DEFAULT;
  if (conf != null) {
    taskCount = conf.getInt(TEZ_SIMPLE_DAG_NUM_TASKS, TEZ_SIMPLE_DAG_NUM_TASKS_DEFAULT);
    payload = TezUtils.createUserPayloadFromConf(conf);
  }
  DAG dag = DAG.create(name);
  Vertex v1 = Vertex.create("v1", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v2 = Vertex.create("v2", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  dag.addVertex(v1).addVertex(v2).addEdge(Edge.create(v1, v2,
      EdgeProperty.create(DataMovementType.SCATTER_GATHER,
          DataSourceType.PERSISTED,
          SchedulingType.SEQUENTIAL,
          TestOutput.getOutputDesc(payload),
          TestInput.getInputDesc(payload))));
  return dag;
}
 

@Test(timeout = 5000)
public void testInvalidConsumerNumber() throws Exception {
  EventHandler mockEventHandler = mock(EventHandler.class);
  Edge edge = new Edge(EdgeProperty.create(
      EdgeManagerPluginDescriptor.create(CustomEdgeManagerWithInvalidReturnValue.class.getName())
        .setUserPayload(new CustomEdgeManagerWithInvalidReturnValue.EdgeManagerConfig(1,1,0,1).toUserPayload()),
      DataSourceType.PERSISTED,
      SchedulingType.SEQUENTIAL,
      OutputDescriptor.create(""),
      InputDescriptor.create("")), mockEventHandler, new TezConfiguration());
  TezVertexID v1Id = createVertexID(1);
  TezVertexID v2Id = createVertexID(2);
  edge.setSourceVertex(mockVertex("v1", v1Id, new LinkedHashMap<TezTaskID, Task>()));
  edge.setDestinationVertex(mockVertex("v2", v2Id, new LinkedHashMap<TezTaskID, Task>()));
  edge.initialize();
  try {
    TezEvent ireEvent = new TezEvent(InputReadErrorEvent.create("diag", 0, 1),
        new EventMetaData(EventProducerConsumerType.INPUT, "v2", "v1",
            TezTaskAttemptID.getInstance(TezTaskID.getInstance(v2Id, 1), 1)));
    edge.sendTezEventToSourceTasks(ireEvent);
    Assert.fail();
  } catch (AMUserCodeException e) {
    e.printStackTrace();
    assertTrue(e.getCause().getMessage().contains("ConsumerTaskNum must be positive"));
  }
}
 

@Test(timeout = 5000)
public void testConvertVertexReconfigreDoneEvent() {
  TezVertexID vId = tezVertexID;
  Map<String, EdgeProperty> edgeMgrs =
      new HashMap<String, EdgeProperty>();

  edgeMgrs.put("a", EdgeProperty.create(EdgeManagerPluginDescriptor.create("a.class")
      .setHistoryText("text"), DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
      OutputDescriptor.create("Out"), InputDescriptor.create("In")));
  VertexConfigurationDoneEvent event = new VertexConfigurationDoneEvent(vId, 0L, 1, null,
      edgeMgrs, null, true);
  HistoryEventProto proto = converter.convert(event);
  assertCommon(proto, HistoryEventType.VERTEX_CONFIGURE_DONE, 0L,
      EntityTypes.TEZ_VERTEX_ID, null, null, 2);
  assertEventData(proto, ATSConstants.NUM_TASKS, "1");
  assertEventData(proto, ATSConstants.UPDATED_EDGE_MANAGERS, null);

  /*
  Map<String, Object> updatedEdgeMgrs = (Map<String, Object>)
      evt.getEventInfo().get(ATSConstants.UPDATED_EDGE_MANAGERS);
  Assert.assertEquals(1, updatedEdgeMgrs.size());
  Assert.assertTrue(updatedEdgeMgrs.containsKey("a"));
  Map<String, Object> updatedEdgeMgr = (Map<String, Object>) updatedEdgeMgrs.get("a");

  Assert.assertEquals(DataMovementType.CUSTOM.name(),
      updatedEdgeMgr.get(DAGUtils.DATA_MOVEMENT_TYPE_KEY));
  Assert.assertEquals("In", updatedEdgeMgr.get(DAGUtils.EDGE_DESTINATION_CLASS_KEY));
  Assert.assertEquals("a.class", updatedEdgeMgr.get(DAGUtils.EDGE_MANAGER_CLASS_KEY));
  */
}
 

@Override
public void setParallelism(int parallelism, VertexLocationHint vertexLocationHint,
    Map<String, EdgeManagerPluginDescriptor> sourceEdgeManagers,
    Map<String, InputSpecUpdate> rootInputSpecUpdates, boolean fromVertexManager) 
        throws AMUserCodeException {
  // temporarily support conversion of edge manager to edge property
  Map<String, EdgeProperty> sourceEdgeProperties = Maps.newHashMap();
  readLock.lock();
  try {
    if (sourceEdgeManagers != null && !sourceEdgeManagers.isEmpty()) {
      for (Edge e : sourceVertices.values()) {
        EdgeManagerPluginDescriptor newEdge = sourceEdgeManagers.get(e.getSourceVertexName());
        EdgeProperty oldEdge = e.getEdgeProperty();
        if (newEdge != null) {
          sourceEdgeProperties.put(
              e.getSourceVertexName(),
              EdgeProperty.create(newEdge, oldEdge.getDataSourceType(),
                  oldEdge.getSchedulingType(), oldEdge.getEdgeSource(),
                  oldEdge.getEdgeDestination()));
        }
      }
    }
  } finally {
    readLock.unlock();
  }
  setParallelismWrapper(parallelism, vertexLocationHint, sourceEdgeProperties, rootInputSpecUpdates,
      fromVertexManager);
}
 

private void reconfigVertex(final int finalTaskParallelism) {
  Map<String, EdgeProperty> edgeProperties =
      new HashMap<String, EdgeProperty>(bipartiteSources);
  Iterable<Map.Entry<String, SourceVertexInfo>> bipartiteItr = getBipartiteInfo();
  for(Map.Entry<String, SourceVertexInfo> entry : bipartiteItr) {
    String vertex = entry.getKey();
    EdgeProperty oldEdgeProp = entry.getValue().edgeProperty;
    EdgeProperty newEdgeProp = EdgeProperty.create(entry.getValue().newDescriptor,
        oldEdgeProp.getDataSourceType(), oldEdgeProp.getSchedulingType(),
        oldEdgeProp.getEdgeSource(), oldEdgeProp.getEdgeDestination());
    edgeProperties.put(vertex, newEdgeProp);
  }
  getContext().reconfigureVertex(finalTaskParallelism, null, edgeProperties);
}
 

private void verifyEdgeProperties(EdgeProperty edgeProperty, String[] sources,
                                  int[] numChunksPerSrc, int maxParallelism)
  throws InvalidProtocolBufferException {
  CartesianProductConfigProto config = CartesianProductConfigProto.parseFrom(ByteString.copyFrom(
    edgeProperty.getEdgeManagerDescriptor().getUserPayload().getPayload()));
  assertArrayEquals(sources, config.getSourcesList().toArray());
  assertArrayEquals(numChunksPerSrc, Ints.toArray(config.getNumChunksList()));
  assertEquals(maxParallelism, config.getMaxParallelism());
}
 

@Override
public synchronized Map<String, EdgeProperty> getOutputVertexEdgeProperties() {
  checkAndThrowIfDone();
  Map<Vertex, Edge> outputs = managedVertex.getOutputVertices();
  Map<String, EdgeProperty> vertexEdgeMap =
                      Maps.newHashMapWithExpectedSize(outputs.size());
  for (Map.Entry<Vertex, Edge> entry : outputs.entrySet()) {
    vertexEdgeMap.put(entry.getKey().getName(), entry.getValue().getEdgeProperty());
  }
  return vertexEdgeMap;
}
 

/**
 * v0 and v1 are two cartesian product sources; v2 is broadcast source
 */
private void setupDAGVertexOnlyWithBroadcast(int maxParallelism, long minWorkloadPerWorker,
                                             int srcParallelismMultiplier) throws Exception {
  Map<String, EdgeProperty> edgePropertyMap = getEdgePropertyMap(2);
  edgePropertyMap.put("v2", EdgeProperty.create(BROADCAST, null, null, null, null));
  when(ctx.getInputVertexEdgeProperties()).thenReturn(edgePropertyMap);
  setSrcParallelism(ctx, srcParallelismMultiplier, 2, 3, 5);

  CartesianProductConfigProto.Builder builder = CartesianProductConfigProto.newBuilder();
  builder.setIsPartitioned(false).addSources("v0").addSources("v1")
    .setMaxParallelism(maxParallelism).setMinOpsPerWorker(minWorkloadPerWorker)
    .setNumPartitionsForFairCase(maxParallelism);
  vertexManager.initialize(builder.build());
}
 

@Override
public void initialize() {
  // this will prevent vertex from starting until we notify we are done
  getContext().vertexReconfigurationPlanned();
  Map<String, EdgeProperty> edges = getContext().getInputVertexEdgeProperties();
  // wait for sources and self to start
  numConfiguredSources = 0;
  configured = new AtomicBoolean(false);
  started = new AtomicBoolean(false);
  for (String entry : edges.keySet()) {
    getContext().registerForVertexStateUpdates(entry, EnumSet.of(VertexState.CONFIGURED));
  }
}
 

private void verifyVertexGroupInfo(EdgeProperty edgeProperty, int positionInGroup,
                                   int... numTaskPerVertexInGroup)
  throws InvalidProtocolBufferException {
  CartesianProductConfigProto config = CartesianProductConfigProto.parseFrom(ByteString.copyFrom(
    edgeProperty.getEdgeManagerDescriptor().getUserPayload().getPayload()));
  assertEquals(positionInGroup, config.getPositionInGroup());
  int i = 0;
  for (int numTask : numTaskPerVertexInGroup) {
    assertEquals(numTask, config.getNumTaskPerVertexInGroup(i));
    i++;
  }
}
 

/**
 * This is a convenience method for creating an Edge descriptor based on the specified
 * EdgeManagerDescriptor.
 *
 * @param edgeManagerDescriptor the custom edge specification
 * @return an {@link org.apache.tez.dag.api.EdgeProperty} instance
 */
public EdgeProperty createDefaultCustomEdgeProperty(EdgeManagerPluginDescriptor edgeManagerDescriptor) {
  Objects.requireNonNull(edgeManagerDescriptor, "EdgeManagerDescriptor cannot be null");
  EdgeProperty edgeProperty =
      EdgeProperty.create(edgeManagerDescriptor, EdgeProperty.DataSourceType.PERSISTED,
          EdgeProperty.SchedulingType.SEQUENTIAL,
          OutputDescriptor.create(getOutputClassName()).setUserPayload(getOutputPayload()),
          InputDescriptor.create(getInputClassName()).setUserPayload(getInputPayload()));
  Utils.setEdgePropertyHistoryText(this, edgeProperty);
  return edgeProperty;
}
 

public void fromProto(VertexConfigurationDoneProto proto) {
  this.vertexID = TezVertexID.fromString(proto.getVertexId());
  this.reconfigureDoneTime = proto.getReconfigureDoneTime();
  this.setParallelismCalledFlag = proto.getSetParallelismCalledFlag();
  this.numTasks = proto.getNumTasks();
  if (proto.hasVertexLocationHint()) {
    this.vertexLocationHint = DagTypeConverters.convertVertexLocationHintFromProto(
        proto.getVertexLocationHint());
  }
  if (proto.getEdgeManagerDescriptorsCount() > 0) {
    this.sourceEdgeProperties = new HashMap<String, EdgeProperty>(
        proto.getEdgeManagerDescriptorsCount());
    for (EdgeManagerDescriptorProto edgeManagerProto :
      proto.getEdgeManagerDescriptorsList()) {
      EdgeProperty edgeProperty =
          DagTypeConverters.convertFromProto(
              edgeManagerProto.getEdgeProperty());
      sourceEdgeProperties.put(edgeManagerProto.getEdgeName(),
          edgeProperty);
    }
  }
  if (proto.getRootInputSpecUpdatesCount() > 0) {
    this.rootInputSpecUpdates = Maps.newHashMap();
    for (RootInputSpecUpdateProto rootInputSpecUpdateProto : proto.getRootInputSpecUpdatesList()) {
      InputSpecUpdate specUpdate;
      if (rootInputSpecUpdateProto.getForAllWorkUnits()) {
        specUpdate = InputSpecUpdate
            .createAllTaskInputSpecUpdate(rootInputSpecUpdateProto.getNumPhysicalInputs(0));
      } else {
        specUpdate = InputSpecUpdate
            .createPerTaskInputSpecUpdate(rootInputSpecUpdateProto.getNumPhysicalInputsList());
      }
      this.rootInputSpecUpdates.put(rootInputSpecUpdateProto.getInputName(), specUpdate);
    }
  }
}
 

/**
 * This is a convenience method for the typical usage of this edge, and creates an instance of
 * {@link org.apache.tez.dag.api.EdgeProperty} which is likely to be used. </p>
 * If custom edge properties are required, the methods to get the relevant payloads should be
 * used. </p>
 * * In this case - DataMovementType.BROADCAST, EdgeProperty.DataSourceType.PERSISTED,
 * EdgeProperty.SchedulingType.SEQUENTIAL
 *
 * @return an {@link org.apache.tez.dag.api.EdgeProperty} instance
 */
public EdgeProperty createDefaultBroadcastEdgeProperty() {
  EdgeProperty edgeProperty = EdgeProperty.create(EdgeProperty.DataMovementType.BROADCAST,
      EdgeProperty.DataSourceType.PERSISTED, EdgeProperty.SchedulingType.SEQUENTIAL,
      OutputDescriptor.create(
          getOutputClassName()).setUserPayload(getOutputPayload()),
      InputDescriptor.create(
          getInputClassName()).setUserPayload(getInputPayload()));
  Utils.setEdgePropertyHistoryText(this, edgeProperty);
  return edgeProperty;
}
 

/**
 * This is a convenience method for creating an Edge descriptor based on the specified
 * EdgeManagerDescriptor.
 *
 * @param edgeManagerDescriptor the custom edge specification
 * @return an {@link org.apache.tez.dag.api.EdgeProperty} instance
 */
public EdgeProperty createDefaultCustomEdgeProperty(EdgeManagerPluginDescriptor edgeManagerDescriptor) {
  Objects.requireNonNull(edgeManagerDescriptor, "EdgeManagerDescriptor cannot be null");
  EdgeProperty edgeProperty =
      EdgeProperty.create(edgeManagerDescriptor, EdgeProperty.DataSourceType.PERSISTED,
          EdgeProperty.SchedulingType.SEQUENTIAL,
          OutputDescriptor.create(getOutputClassName()).setUserPayload(getOutputPayload()),
          InputDescriptor.create(getInputClassName()).setUserPayload(getInputPayload()));
  Utils.setEdgePropertyHistoryText(this, edgeProperty);
  return edgeProperty;
}
 

@Test(timeout = 5000)
public void testDAGVertexOnlyGroupByMaxParallelism() throws Exception {
  setupDAGVertexOnly(30, 1, 30, 1);
  vertexManager.onVertexStateUpdated(new VertexStateUpdate("v0", VertexState.CONFIGURED));
  vertexManager.onVertexStateUpdated(new VertexStateUpdate("v1", VertexState.CONFIGURED));

  vertexManager.onVertexManagerEventReceived(getVMEvent(250, "v0", 0));
  verify(ctx, never()).reconfigureVertex(
    anyInt(), any(VertexLocationHint.class), anyMapOf(String.class, EdgeProperty.class));

  vertexManager.onVertexManagerEventReceived(getVMEvent(200, "v1", 0));
  verify(ctx, times(1)).reconfigureVertex(
    eq(30), any(VertexLocationHint.class), edgePropertiesCaptor.capture());
  Map<String, EdgeProperty> edgeProperties = edgePropertiesCaptor.getValue();
  verifyEdgeProperties(edgeProperties.get("v0"), new String[]{"v0", "v1"}, new int[]{5, 6}, 30);
  verifyVertexGroupInfo(edgeProperties.get("v0"), 0);
  verifyEdgeProperties(edgeProperties.get("v1"), new String[]{"v0", "v1"}, new int[]{5, 6}, 30);
  verifyVertexGroupInfo(edgeProperties.get("v1"), 0);

  vertexManager.onVertexStarted(null);
  verifyScheduleRequest(0);
  vertexManager.onSourceTaskCompleted(getTaId("v0", 0));
  vertexManager.onSourceTaskCompleted(getTaId("v1", 0));
  verifyScheduleRequest(1, 0, 6, 1, 7);
  vertexManager.onSourceTaskCompleted(getTaId("v0", 1));
  verifyScheduleRequest(2, 12, 13, 18, 19, 24, 25);
}
 

private HistoryEventProto convertVertexReconfigureDoneEvent(VertexConfigurationDoneEvent event) {
  HistoryEventProto.Builder builder = makeBuilderForEvent(event, event.getReconfigureDoneTime(),
      null, null, null, event.getVertexID(), null, null, null);
  if (event.getSourceEdgeProperties() != null && !event.getSourceEdgeProperties().isEmpty()) {
    Map<String, Object> updatedEdgeManagers = new HashMap<String, Object>();
    for (Entry<String, EdgeProperty> entry : event.getSourceEdgeProperties().entrySet()) {
      updatedEdgeManagers.put(entry.getKey(), DAGUtils.convertEdgeProperty(entry.getValue()));
    }
    addEventData(builder, ATSConstants.UPDATED_EDGE_MANAGERS, updatedEdgeManagers);
  }
  addEventData(builder, ATSConstants.NUM_TASKS, event.getNumTasks());
  return builder.build();
}
 

/**
 * create a DAG with 2 vertices (v1 --> v2), set payload on Input/Output/Processor/VertexManagerPlugin to
 * control where throw exception
 * 
 * @param exLocation
 * @return
 * @throws IOException
 */
private DAG createDAG(ExceptionLocation exLocation) throws IOException {
  DAG dag = DAG.create("dag_" + exLocation.name());
  UserPayload payload =
      UserPayload.create(ByteBuffer.wrap(exLocation.name().getBytes()));
  Vertex v1 =
      Vertex.create("v1", ProcessorWithException.getProcDesc(payload), 1);
  InputDescriptor inputDesc = InputWithException.getInputDesc(payload);
  InputInitializerDescriptor iiDesc =
      InputInitializerWithException.getIIDesc(payload);
  v1.addDataSource("input",
      DataSourceDescriptor.create(inputDesc, iiDesc, null));
  v1.setVertexManagerPlugin(RootInputVertexManagerWithException
      .getVMDesc(exLocation));

  Vertex v2 = 
      Vertex.create("v2", DoNothingProcessor.getProcDesc(), 1);
  v2.addDataSource("input2",
      DataSourceDescriptor.create(InputDescriptor.create(NoOpInput.class.getName()),
        InputInitializerWithException2.getIIDesc(payload), null));

  dag.addVertex(v1)
    .addVertex(v2);
  if (exLocation.name().startsWith("EM_")) {
    dag.addEdge(Edge.create(v1, v2, EdgeProperty.create(
        EdgeManagerPluginDescriptor.create(CustomEdgeManager.class.getName())
          .setUserPayload(payload),
        DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
        OutputWithException.getOutputDesc(payload), InputWithException.getInputDesc(payload))));
  } else {
    // set Customized VertexManager here, it can't been used for CustomEdge
    v2.setVertexManagerPlugin(InputReadyVertexManagerWithException.getVMDesc(exLocation));
    dag.addEdge(Edge.create(v1, v2, EdgeProperty.create(DataMovementType.ONE_TO_ONE,
        DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
        OutputWithException.getOutputDesc(payload), InputWithException.getInputDesc(payload))));
  }

  return dag;
}
 

void logVertexConfigurationDoneEvent() {
  if (recoveryData == null || !recoveryData.shouldSkipInit()) {
    Map<String, EdgeProperty> sourceEdgeProperties = new HashMap<String, EdgeProperty>();
    for (Map.Entry<Vertex, Edge> entry : this.sourceVertices.entrySet()) {
      sourceEdgeProperties.put(entry.getKey().getName(), entry.getValue().getEdgeProperty());
    }
    VertexConfigurationDoneEvent reconfigureDoneEvent =
        new VertexConfigurationDoneEvent(vertexId, clock.getTime(),
            numTasks, taskLocationHints == null ? null : VertexLocationHint.create(Lists.newArrayList(taskLocationHints)),
                sourceEdgeProperties, rootInputSpecs, setParallelismCalledFlag);
    this.appContext.getHistoryHandler().handle(
        new DAGHistoryEvent(getDAGId(), reconfigureDoneEvent));
  }
}
 

public static DAG createDAG(String name, 
    Configuration conf) throws Exception {
  UserPayload payload = UserPayload.create(null);
  int taskCount = TEZ_SIMPLE_V_DAG_NUM_TASKS_DEFAULT;
  if (conf != null) {
    taskCount = conf.getInt(TEZ_SIMPLE_V_DAG_NUM_TASKS, TEZ_SIMPLE_V_DAG_NUM_TASKS_DEFAULT);
    payload = TezUtils.createUserPayloadFromConf(conf);
  }
  DAG dag = DAG.create(name);
  Vertex v1 = Vertex.create("v1", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v2 = Vertex.create("v2", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v3 = Vertex.create("v3", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  dag.addVertex(v1).addVertex(v2).addVertex(v3);
  dag.addEdge(Edge.create(v1, v3,
      EdgeProperty.create(DataMovementType.SCATTER_GATHER,
          DataSourceType.PERSISTED,
          SchedulingType.SEQUENTIAL,
          TestOutput.getOutputDesc(payload),
          TestInput.getInputDesc(payload))));
  dag.addEdge(Edge.create(v2, v3,
      EdgeProperty.create(DataMovementType.SCATTER_GATHER,
          DataSourceType.PERSISTED,
          SchedulingType.SEQUENTIAL,
          TestOutput.getOutputDesc(payload),
          TestInput.getInputDesc(payload))));
  return dag;
}
 

/**
 * Adds an edge to given 2 vertices.
 * @param v1 vertice 1
 * @param v2 vertice 2
 * @param dataMovementType Data movement type
 */
protected static void addEdge(Vertex v1, Vertex v2, DataMovementType dataMovementType) {
    dag.addEdge(Edge.create(v1, v2,
        EdgeProperty.create(dataMovementType,
            DataSourceType.PERSISTED,
            SchedulingType.SEQUENTIAL,
            TestOutput.getOutputDesc(payload),
            TestInput.getInputDesc(payload))));
}
 

public static DAG createDAG(String name, 
    Configuration conf) throws Exception {
  UserPayload payload = UserPayload.create(null);
  int taskCount = TEZ_SIMPLE_REVERSE_V_DAG_NUM_TASKS_DEFAULT;
  if (conf != null) {
    taskCount = conf.getInt(TEZ_SIMPLE_REVERSE_V_DAG_NUM_TASKS, TEZ_SIMPLE_REVERSE_V_DAG_NUM_TASKS_DEFAULT);
    payload = TezUtils.createUserPayloadFromConf(conf);
  }
  DAG dag = DAG.create(name);
  Vertex v1 = Vertex.create("v1", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v2 = Vertex.create("v2", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v3 = Vertex.create("v3", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  dag.addVertex(v1).addVertex(v2).addVertex(v3);
  dag.addEdge(Edge.create(v1, v2,
      EdgeProperty.create(DataMovementType.SCATTER_GATHER,
          DataSourceType.PERSISTED,
          SchedulingType.SEQUENTIAL,
          TestOutput.getOutputDesc(payload),
          TestInput.getInputDesc(payload))));
  dag.addEdge(Edge.create(v1, v3,
      EdgeProperty.create(DataMovementType.SCATTER_GATHER,
          DataSourceType.PERSISTED,
          SchedulingType.SEQUENTIAL,
          TestOutput.getOutputDesc(payload),
          TestInput.getInputDesc(payload))));
  return dag;
}
 

/**
 * This is a convenience method for the typical usage of this edge, and creates an instance of
 * {@link org.apache.tez.dag.api.EdgeProperty} which is likely to be used. </p>
 * * In this case - DataMovementType.SCATTER_GATHER, EdgeProperty.DataSourceType.PERSISTED,
 * EdgeProperty.SchedulingType.SEQUENTIAL
 *
 * @return an {@link org.apache.tez.dag.api.EdgeProperty} instance
 */
public EdgeProperty createDefaultEdgeProperty() {
  EdgeProperty edgeProperty = EdgeProperty.create(EdgeProperty.DataMovementType.SCATTER_GATHER,
      EdgeProperty.DataSourceType.PERSISTED, EdgeProperty.SchedulingType.SEQUENTIAL,
      OutputDescriptor.create(
          getOutputClassName()).setUserPayload(getOutputPayload()),
      InputDescriptor.create(
          getInputClassName()).setUserPayload(getInputPayload()));
  Utils.setEdgePropertyHistoryText(this, edgeProperty);
  return edgeProperty;
}