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

/**
 * v1 --> v2 <br>
 * v1 has a customized VM to control whether to schedule only one second task when it is partiallyFinished test case.
 * v2 has a customized VM which could control when to kill AM
 *
 * @param vertexManagerClass
 * @param dmType
 * @param failOnParitialCompleted
 * @return
 * @throws IOException
 */
private DAG createDAG(String dagName, Class vertexManagerClass, DataMovementType dmType,
    boolean failOnParitialCompleted) throws IOException {
  if (failOnParitialCompleted) {
    tezConf.set(FAIL_ON_PARTIAL_FINISHED, "true");
  } else {
    tezConf.set(FAIL_ON_PARTIAL_FINISHED, "false");
  }
  DAG dag = DAG.create(dagName);
  UserPayload payload = UserPayload.create(null);
  Vertex v1 = Vertex.create("v1", MyProcessor.getProcDesc(), 2);
  v1.setVertexManagerPlugin(VertexManagerPluginDescriptor.create(
      ScheduleControlledVertexManager.class.getName()).setUserPayload(
      TezUtils.createUserPayloadFromConf(tezConf)));
  Vertex v2 = Vertex.create("v2", DoNothingProcessor.getProcDesc(), 2);
  v2.setVertexManagerPlugin(VertexManagerPluginDescriptor.create(
      vertexManagerClass.getName()).setUserPayload(
      TezUtils.createUserPayloadFromConf(tezConf)));

  dag.addVertex(v1).addVertex(v2);
  dag.addEdge(Edge.create(v1, v2, EdgeProperty.create(dmType,
      DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
      TestOutput.getOutputDesc(payload), TestInput.getInputDesc(payload))));
  return dag;
}
 

public static DAG createDAG(String name, 
    Configuration conf) throws Exception {
  byte[] payload = 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 = new DAG(name);
  Vertex v1 = new Vertex("v1", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v2 = new Vertex("v2", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  dag.addVertex(v1).addVertex(v2).addEdge(new Edge(v1, v2, 
      new EdgeProperty(DataMovementType.SCATTER_GATHER, 
          DataSourceType.PERSISTED, 
          SchedulingType.SEQUENTIAL, 
          TestOutput.getOutputDesc(payload), 
          TestInput.getInputDesc(payload))));
  return dag;
}
 

public static DAG createDAG(String name, 
    Configuration conf) throws Exception {
  UserPayload payload = 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);
  Vertex v3 = Vertex.create("v3", 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))));
  dag.addVertex(v3).addEdge(Edge.create(v2, v3,
      EdgeProperty.create(DataMovementType.SCATTER_GATHER,
          DataSourceType.PERSISTED,
          SchedulingType.SEQUENTIAL,
          TestOutput.getOutputDesc(payload),
          TestInput.getInputDesc(payload))));
  return dag;
}
 

private DAG createDAG(int numGenTasks, int totalSourceDataSize, int numFetcherTasks) {
  int bytesPerSource = totalSourceDataSize / numGenTasks;
  LOG.info("DataPerSourceTask(bytes)=" + bytesPerSource);
  ByteBuffer payload = ByteBuffer.allocate(4);
  payload.putInt(0, bytesPerSource);

  Vertex broadcastVertex = Vertex.create("DataGen",
      ProcessorDescriptor.create(InputGenProcessor.class.getName())
          .setUserPayload(UserPayload.create(payload)), numGenTasks);
  Vertex fetchVertex = Vertex.create("FetchVertex",
      ProcessorDescriptor.create(InputFetchProcessor.class.getName()), numFetcherTasks);
  UnorderedKVEdgeConfig edgeConf = UnorderedKVEdgeConfig.newBuilder(NullWritable.class
  .getName(), IntWritable.class.getName()).setCompression(false, null, null).build();

  DAG dag = DAG.create("BroadcastLoadGen");
  dag.addVertex(broadcastVertex).addVertex(fetchVertex).addEdge(
      Edge.create(broadcastVertex, fetchVertex, edgeConf.createDefaultBroadcastEdgeProperty()));
  return dag;
}
 

public static DAG createDAG(String name, 
    Configuration conf) throws Exception {
  byte[] payload = 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 = new DAG(name);
  Vertex v1 = new Vertex("v1", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v2 = new Vertex("v2", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v3 = new Vertex("v3", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  dag.addVertex(v1).addVertex(v2).addEdge(new Edge(v1, v2, 
      new EdgeProperty(DataMovementType.SCATTER_GATHER, 
          DataSourceType.PERSISTED, 
          SchedulingType.SEQUENTIAL, 
          TestOutput.getOutputDesc(payload), 
          TestInput.getInputDesc(payload))));
  dag.addVertex(v3).addEdge(new Edge(v2, v3, 
          new EdgeProperty(DataMovementType.SCATTER_GATHER, 
              DataSourceType.PERSISTED, 
              SchedulingType.SEQUENTIAL, 
              TestOutput.getOutputDesc(payload), 
              TestInput.getInputDesc(payload))));
  return dag;
}
 

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

private DAG createDAG(TezConfiguration tezConf) throws IOException {
  InputDescriptor inputDescriptor = InputDescriptor.create(FakeInput.class.getName());
  InputInitializerDescriptor inputInitializerDescriptor =
    InputInitializerDescriptor.create(FakeInputInitializer.class.getName());
  DataSourceDescriptor dataSourceDescriptor =
    DataSourceDescriptor.create(inputDescriptor, inputInitializerDescriptor, null);

  Vertex v1 = Vertex.create(VERTEX1, ProcessorDescriptor.create(TokenProcessor.class.getName()));
  v1.addDataSource(INPUT, dataSourceDescriptor);
  Vertex v2 = Vertex.create(VERTEX2, ProcessorDescriptor.create(TokenProcessor.class.getName()));
  v2.addDataSource(INPUT, dataSourceDescriptor);

  OutputDescriptor outputDescriptor = OutputDescriptor.create(FakeOutput.class.getName());
  OutputCommitterDescriptor outputCommitterDescriptor =
    OutputCommitterDescriptor.create(FakeOutputCommitter.class.getName());
  DataSinkDescriptor dataSinkDescriptor =
    DataSinkDescriptor.create(outputDescriptor, outputCommitterDescriptor, null);

  CartesianProductConfig cartesianProductConfig =
    new CartesianProductConfig(Arrays.asList(sourceVertices));
  UserPayload userPayload = cartesianProductConfig.toUserPayload(tezConf);

  Vertex v3 = Vertex.create(VERTEX3, ProcessorDescriptor.create(JoinProcessor.class.getName()));
  v3.addDataSink(OUTPUT, dataSinkDescriptor);
  v3.setVertexManagerPlugin(
    VertexManagerPluginDescriptor.create(CartesianProductVertexManager.class.getName())
                                 .setUserPayload(userPayload));

  EdgeManagerPluginDescriptor edgeManagerDescriptor =
    EdgeManagerPluginDescriptor.create(CartesianProductEdgeManager.class.getName());
  edgeManagerDescriptor.setUserPayload(userPayload);
  UnorderedPartitionedKVEdgeConfig edgeConf =
    UnorderedPartitionedKVEdgeConfig.newBuilder(Text.class.getName(), IntWritable.class.getName(),
      RoundRobinPartitioner.class.getName()).build();
  EdgeProperty edgeProperty = edgeConf.createDefaultCustomEdgeProperty(edgeManagerDescriptor);

  return DAG.create("CrossProduct").addVertex(v1).addVertex(v2).addVertex(v3)
    .addEdge(Edge.create(v1, v3, edgeProperty)).addEdge(Edge.create(v2, v3, edgeProperty));
}
 

@Test (timeout=60000)
public void testBasicSuccessBroadcast() throws Exception {
  DAG dag = new DAG("testBasicSuccessBroadcast");
  Vertex v1 = new Vertex("v1", TestProcessor.getProcDesc(null), 2, SimpleTestDAG.defaultResource);
  Vertex v2 = new Vertex("v2", TestProcessor.getProcDesc(null), 2, SimpleTestDAG.defaultResource);
  dag.addVertex(v1).addVertex(v2).addEdge(new Edge(v1, v2, 
      new EdgeProperty(DataMovementType.BROADCAST, 
          DataSourceType.PERSISTED, 
          SchedulingType.SEQUENTIAL, 
          TestOutput.getOutputDesc(null), 
          TestInput.getInputDesc(null))));
  runDAGAndVerify(dag, DAGStatus.State.SUCCEEDED);
}
 

@Ignore
@Test (timeout = 600000)
public void testMemoryBroadcast() throws Exception {
  DAG dag = DAG.create("testMemoryBroadcast");
  Vertex vA = Vertex.create("A", ProcessorDescriptor.create("Proc.class"), numTasks);
  Vertex vB = Vertex.create("B", ProcessorDescriptor.create("Proc.class"), numTasks);
  dag.addVertex(vA)
      .addVertex(vB)
      .addEdge(
          Edge.create(vA, vB, EdgeProperty.create(DataMovementType.BROADCAST,
              DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
              OutputDescriptor.create("Out"), InputDescriptor.create("In"))));
  testMemory(dag, true);
}
 

@Ignore
@Test (timeout = 600000)
public void testMemoryOneToOne() throws Exception {
  DAG dag = DAG.create("testMemoryOneToOne");
  Vertex vA = Vertex.create("A", ProcessorDescriptor.create("Proc.class"), numTasks);
  Vertex vB = Vertex.create("B", ProcessorDescriptor.create("Proc.class"), numTasks);
  dag.addVertex(vA)
      .addVertex(vB)
      .addEdge(
          Edge.create(vA, vB, EdgeProperty.create(DataMovementType.ONE_TO_ONE,
              DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
              OutputDescriptor.create("Out"), InputDescriptor.create("In"))));
  testMemory(dag, true);
}
 

@Test (timeout=60000)
public void testBasicSuccessBroadcast() throws Exception {
  DAG dag = DAG.create("testBasicSuccessBroadcast");
  Vertex v1 =
      Vertex.create("v1", TestProcessor.getProcDesc(null), 2, SimpleTestDAG.defaultResource);
  Vertex v2 =
      Vertex.create("v2", TestProcessor.getProcDesc(null), 2, SimpleTestDAG.defaultResource);
  dag.addVertex(v1).addVertex(v2).addEdge(Edge.create(v1, v2,
      EdgeProperty.create(DataMovementType.BROADCAST,
          DataSourceType.PERSISTED,
          SchedulingType.SEQUENTIAL,
          TestOutput.getOutputDesc(null),
          TestInput.getInputDesc(null))));
  runDAGAndVerify(dag, DAGStatus.State.SUCCEEDED);
}
 

@Test(timeout = 60000)
public void testVertexFailuresMaxPercent() throws TezException, InterruptedException, IOException {

  TezConfiguration tezConf = new TezConfiguration(mrrTezCluster.getConfig());
  tezConf.set(TezConfiguration.TEZ_VERTEX_FAILURES_MAXPERCENT, "50.0f");
  tezConf.setInt(TezConfiguration.TEZ_AM_TASK_MAX_FAILED_ATTEMPTS, 1);
  TezClient tezClient = TezClient.create("TestVertexFailuresMaxPercent", tezConf);
  tezClient.start();

  try {
    DAG dag = DAG.create("TestVertexFailuresMaxPercent");
    Vertex vertex1 = Vertex.create("Parent", ProcessorDescriptor.create(
        FailingAttemptProcessor.class.getName()), 2);
    Vertex vertex2 = Vertex.create("Child", ProcessorDescriptor.create(FailingAttemptProcessor.class.getName()), 2);

    OrderedPartitionedKVEdgeConfig edgeConfig = OrderedPartitionedKVEdgeConfig
        .newBuilder(Text.class.getName(), IntWritable.class.getName(),
            HashPartitioner.class.getName())
        .setFromConfiguration(tezConf)
        .build();
    dag.addVertex(vertex1)
        .addVertex(vertex2)
        .addEdge(Edge.create(vertex1, vertex2, edgeConfig.createDefaultEdgeProperty()));

    DAGClient dagClient = tezClient.submitDAG(dag);
    dagClient.waitForCompletion();
    Assert.assertEquals(DAGStatus.State.SUCCEEDED, dagClient.getDAGStatus(null).getState());
  } finally {
    tezClient.stop();
  }
}
 

/**
 * 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;
}
 

public static DAG createDAG(String name,
    Configuration conf) throws Exception {
  UserPayload payload = UserPayload.create(null);
  int taskCount = MULTI_ATTEMPT_DAG_VERTEX_NUM_TASKS_DEFAULT;
  if (conf != null) {
    taskCount = conf.getInt(MULTI_ATTEMPT_DAG_VERTEX_NUM_TASKS, MULTI_ATTEMPT_DAG_VERTEX_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);

  // Make each vertex manager fail on appropriate attempt
  v1.setVertexManagerPlugin(VertexManagerPluginDescriptor.create(
      FailOnAttemptVertexManagerPlugin.class.getName())
      .setUserPayload(UserPayload.create(ByteBuffer.wrap(new String("1").getBytes()))));
  v2.setVertexManagerPlugin(VertexManagerPluginDescriptor.create(
      FailOnAttemptVertexManagerPlugin.class.getName())
      .setUserPayload(UserPayload.create(ByteBuffer.wrap(new String("2").getBytes()))));
  v3.setVertexManagerPlugin(VertexManagerPluginDescriptor.create(
      FailOnAttemptVertexManagerPlugin.class.getName())
      .setUserPayload(UserPayload.create(ByteBuffer.wrap(new String("3").getBytes()))));
  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(v2, v3,
      EdgeProperty.create(DataMovementType.SCATTER_GATHER,
          DataSourceType.PERSISTED,
          SchedulingType.SEQUENTIAL,
          TestOutput.getOutputDesc(payload),
          TestInput.getInputDesc(payload))));
  return dag;
}
 

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

/**
 * 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_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;
}
 

@Ignore
@Test (timeout = 600000)
public void testMemoryScatterGather() throws Exception {
  DAG dag = DAG.create("testMemoryScatterGather");
  Vertex vA = Vertex.create("A", ProcessorDescriptor.create("Proc.class"), numTasks);
  Vertex vB = Vertex.create("B", ProcessorDescriptor.create("Proc.class"), numTasks);
  dag.addVertex(vA)
      .addVertex(vB)
      .addEdge(
          Edge.create(vA, vB, EdgeProperty.create(DataMovementType.SCATTER_GATHER,
              DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
              OutputDescriptor.create("Out"), InputDescriptor.create("In"))));
  testMemory(dag, true);
}
 

DAG createDAG(DataMovementType dmType) {
  DAG dag = DAG.create("test-" + dagCount++);
  Vertex vA = Vertex.create("A", ProcessorDescriptor.create("Proc.class"), 5);
  Vertex vB = Vertex.create("B", ProcessorDescriptor.create("Proc.class"), 5);
  Edge eAB = Edge.create(vA, vB, 
  EdgeProperty.create(dmType, DataSourceType.PERSISTED,
      SchedulingType.SEQUENTIAL, OutputDescriptor.create("O.class"),
      InputDescriptor.create("I.class")));
  
  dag.addVertex(vA).addVertex(vB).addEdge(eAB);
  return dag;
}
 

private DAG createDAG(FileSystem fs, TezConfiguration tezConf,
    Map<String, LocalResource> localResources, Path stagingDir,
    String inputPath, String outputPath) throws IOException {

  Configuration inputConf = new Configuration(tezConf);
  inputConf.set(FileInputFormat.INPUT_DIR, inputPath);
  InputDescriptor id = new InputDescriptor(MRInput.class.getName())
      .setUserPayload(MRInput.createUserPayload(inputConf,
          TextInputFormat.class.getName(), true, true));

  Configuration outputConf = new Configuration(tezConf);
  outputConf.set(FileOutputFormat.OUTDIR, outputPath);
  OutputDescriptor od = new OutputDescriptor(MROutput.class.getName())
    .setUserPayload(MROutput.createUserPayload(
        outputConf, TextOutputFormat.class.getName(), true));

  Vertex tokenizerVertex = new Vertex("tokenizer", new ProcessorDescriptor(
      TokenProcessor.class.getName()), -1, MRHelpers.getMapResource(tezConf));
  tokenizerVertex.addInput("MRInput", id, MRInputAMSplitGenerator.class);

  Vertex summerVertex = new Vertex("summer",
      new ProcessorDescriptor(
          SumProcessor.class.getName()), 1, MRHelpers.getReduceResource(tezConf));
  summerVertex.addOutput("MROutput", od, MROutputCommitter.class);

  OrderedPartitionedKVEdgeConfigurer edgeConf = OrderedPartitionedKVEdgeConfigurer
      .newBuilder(Text.class.getName(), IntWritable.class.getName(),
          HashPartitioner.class.getName(), null).build();

  DAG dag = new DAG("WordCount");
  dag.addVertex(tokenizerVertex)
      .addVertex(summerVertex)
      .addEdge(
          new Edge(tokenizerVertex, summerVertex, edgeConf.createDefaultEdgeProperty()));
  return dag;  
}
 

public static DAG createDAG(String name, 
    Configuration conf) throws Exception {
  byte[] payload = 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 = new DAG(name);
  Vertex v1 = new Vertex("v1", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v2 = new Vertex("v2", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v3 = new Vertex("v3", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  dag.addVertex(v1).addVertex(v2).addVertex(v3);
  dag.addEdge(new Edge(v1, v3, 
      new EdgeProperty(DataMovementType.SCATTER_GATHER, 
          DataSourceType.PERSISTED, 
          SchedulingType.SEQUENTIAL, 
          TestOutput.getOutputDesc(payload), 
          TestInput.getInputDesc(payload))));
  dag.addEdge(new Edge(v2, v3, 
          new EdgeProperty(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(new Edge(v1, v2, 
        new EdgeProperty(dataMovementType, 
            DataSourceType.PERSISTED, 
            SchedulingType.SEQUENTIAL, 
            TestOutput.getOutputDesc(payload), 
            TestInput.getInputDesc(payload))));
}
 

public static DAG createDAG(String name, 
    Configuration conf) throws Exception {
  byte[] payload = 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 = new DAG(name);
  Vertex v1 = new Vertex("v1", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v2 = new Vertex("v2", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v3 = new Vertex("v3", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  dag.addVertex(v1).addVertex(v2).addVertex(v3);
  dag.addEdge(new Edge(v1, v2, 
      new EdgeProperty(DataMovementType.SCATTER_GATHER, 
          DataSourceType.PERSISTED, 
          SchedulingType.SEQUENTIAL, 
          TestOutput.getOutputDesc(payload), 
          TestInput.getInputDesc(payload))));
  dag.addEdge(new Edge(v1, v3, 
          new EdgeProperty(DataMovementType.SCATTER_GATHER, 
              DataSourceType.PERSISTED, 
              SchedulingType.SEQUENTIAL, 
              TestOutput.getOutputDesc(payload), 
              TestInput.getInputDesc(payload))));
  return dag;
}
 

public static DAG createDAG(String name,
    Configuration conf) throws Exception {
  byte[] payload = null;
  int taskCount = MULTI_ATTEMPT_DAG_VERTEX_NUM_TASKS_DEFAULT;
  if (conf != null) {
    taskCount = conf.getInt(MULTI_ATTEMPT_DAG_VERTEX_NUM_TASKS, MULTI_ATTEMPT_DAG_VERTEX_NUM_TASKS_DEFAULT);
    payload = TezUtils.createUserPayloadFromConf(conf);
  }
  DAG dag = new DAG(name);
  Vertex v1 = new Vertex("v1", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v2 = new Vertex("v2", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v3 = new Vertex("v3", TestProcessor.getProcDesc(payload), taskCount, defaultResource);

  // Make each vertex manager fail on appropriate attempt
  v1.setVertexManagerPlugin(new VertexManagerPluginDescriptor(
      FailOnAttemptVertexManagerPlugin.class.getName())
      .setUserPayload(new String("1").getBytes()));
  v2.setVertexManagerPlugin(new VertexManagerPluginDescriptor(
      FailOnAttemptVertexManagerPlugin.class.getName())
      .setUserPayload(new String("2").getBytes()));
  v3.setVertexManagerPlugin(new VertexManagerPluginDescriptor(
      FailOnAttemptVertexManagerPlugin.class.getName())
      .setUserPayload(new String("3").getBytes()));
  dag.addVertex(v1).addVertex(v2).addVertex(v3);
  dag.addEdge(new Edge(v1, v2,
      new EdgeProperty(DataMovementType.SCATTER_GATHER,
          DataSourceType.PERSISTED,
          SchedulingType.SEQUENTIAL,
          TestOutput.getOutputDesc(payload),
          TestInput.getInputDesc(payload))));
  dag.addEdge(new Edge(v2, v3,
      new EdgeProperty(DataMovementType.SCATTER_GATHER,
          DataSourceType.PERSISTED,
          SchedulingType.SEQUENTIAL,
          TestOutput.getOutputDesc(payload),
          TestInput.getInputDesc(payload))));
  return dag;
}
 

/**
 * Test basic speculation per vertex conf.
 *
 * @throws Exception the exception
 */
@Retry
@Test (timeout=10000)
public void testBasicSpeculationPerVertexConf() throws Exception {
  DAG dag = DAG.create("test");
  String vNameNoSpec = "A";
  String vNameSpec = "B";
  String speculatorSleepTime = "50";
  Vertex vA = Vertex.create(vNameNoSpec, ProcessorDescriptor.create("Proc.class"), 5);
  Vertex vB = Vertex.create(vNameSpec, ProcessorDescriptor.create("Proc.class"), 5);
  vA.setConf(TezConfiguration.TEZ_AM_SPECULATION_ENABLED, "false");
  dag.addVertex(vA);
  dag.addVertex(vB);
  // min/max src fraction is set to 1. So vertices will run sequentially
  dag.addEdge(
      Edge.create(vA, vB,
          EdgeProperty.create(DataMovementType.SCATTER_GATHER, DataSourceType.PERSISTED,
              SchedulingType.SEQUENTIAL, OutputDescriptor.create("O"),
              InputDescriptor.create("I"))));

  MockTezClient tezClient = createTezSession();

  DAGClient dagClient = tezClient.submitDAG(dag);
  DAGImpl dagImpl = (DAGImpl) mockApp.getContext().getCurrentDAG();
  TezVertexID vertexId = dagImpl.getVertex(vNameSpec).getVertexId();
  TezVertexID vertexIdNoSpec = dagImpl.getVertex(vNameNoSpec).getVertexId();
  // original attempt is killed and speculative one is successful
  TezTaskAttemptID killedTaId =
      TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexId, 0), 0);
  TezTaskAttemptID noSpecTaId = TezTaskAttemptID
      .getInstance(TezTaskID.getInstance(vertexIdNoSpec, 0), 0);

  // cause speculation trigger for both
  mockLauncher.setStatusUpdatesForTask(killedTaId, 100);
  mockLauncher.setStatusUpdatesForTask(noSpecTaId, 100);

  mockLauncher.startScheduling(true);
  org.apache.tez.dag.app.dag.Vertex vSpec = dagImpl.getVertex(vertexId);
  org.apache.tez.dag.app.dag.Vertex vNoSpec = dagImpl.getVertex(vertexIdNoSpec);
  // Wait enough time to give chance for the speculator to trigger
  // speculation on VB.
  // This would fail because of JUnit time out.
  do {
    Thread.sleep(100);
  } while (vSpec.getAllCounters().findCounter(TaskCounter.NUM_SPECULATIONS)
      .getValue() <= 0);
  dagClient.waitForCompletion();
  // speculation for vA but not for vB
  Assert.assertTrue("Num Speculations is not higher than 0",
      vSpec.getAllCounters().findCounter(TaskCounter.NUM_SPECULATIONS)
          .getValue() > 0);
  Assert.assertEquals(0,
      vNoSpec.getAllCounters().findCounter(TaskCounter.NUM_SPECULATIONS)
          .getValue());

  tezClient.stop();
}
 

@Test (timeout = 5000)
public void testBasicEvents() throws Exception {
  TezConfiguration tezconf = new TezConfiguration(defaultConf);
  
  MockTezClient tezClient = new MockTezClient("testMockAM", tezconf, true, null, null, null, null);
  tezClient.start();
  
  MockDAGAppMaster mockApp = tezClient.getLocalClient().getMockApp();
  MockContainerLauncher mockLauncher = mockApp.getContainerLauncher();
  mockLauncher.startScheduling(false);
  mockApp.eventsDelegate = new TestEventsDelegate();
  DAG dag = DAG.create("testBasicEvents");
  Vertex vA = Vertex.create("A", ProcessorDescriptor.create("Proc.class"), 2);
  Vertex vB = Vertex.create("B", ProcessorDescriptor.create("Proc.class"), 2);
  Vertex vC = Vertex.create("C", ProcessorDescriptor.create("Proc.class"), 2);
  Vertex vD = Vertex.create("D", ProcessorDescriptor.create("Proc.class"), 2);
  dag.addVertex(vA)
      .addVertex(vB)
      .addVertex(vC)
      .addVertex(vD)
      .addEdge(
          Edge.create(vA, vB, EdgeProperty.create(DataMovementType.BROADCAST,
              DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
              OutputDescriptor.create("Out"), InputDescriptor.create("In"))))
      .addEdge(
          Edge.create(vA, vC, EdgeProperty.create(DataMovementType.SCATTER_GATHER,
              DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
              OutputDescriptor.create("Out"), InputDescriptor.create("In"))))
      .addEdge(
          Edge.create(vA, vD, EdgeProperty.create(DataMovementType.ONE_TO_ONE,
              DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
              OutputDescriptor.create("Out"), InputDescriptor.create("In"))));

  DAGClient dagClient = tezClient.submitDAG(dag);
  mockLauncher.waitTillContainersLaunched();
  DAGImpl dagImpl = (DAGImpl) mockApp.getContext().getCurrentDAG();
  mockLauncher.startScheduling(true);
  dagClient.waitForCompletion();
  Assert.assertEquals(DAGStatus.State.SUCCEEDED, dagClient.getDAGStatus(null).getState());
  VertexImpl vImpl = (VertexImpl) dagImpl.getVertex(vB.getName());
  TaskImpl tImpl = (TaskImpl) vImpl.getTask(1);
  TezTaskAttemptID taId = TezTaskAttemptID.getInstance(tImpl.getTaskId(), 0);
  List<TezEvent> tEvents = vImpl.getTaskAttemptTezEvents(taId, 0, 0, 1000).getEvents();
  Assert.assertEquals(2, tEvents.size()); // 2 from vA
  Assert.assertEquals(vA.getName(), tEvents.get(0).getDestinationInfo().getEdgeVertexName());
  Assert.assertEquals(0, ((DataMovementEvent)tEvents.get(0).getEvent()).getSourceIndex());
  Assert.assertEquals(vA.getName(), tEvents.get(1).getDestinationInfo().getEdgeVertexName());
  Assert.assertEquals(0, ((DataMovementEvent)tEvents.get(1).getEvent()).getSourceIndex());
  int targetIndex1 = ((DataMovementEvent)tEvents.get(0).getEvent()).getTargetIndex();
  int targetIndex2 = ((DataMovementEvent)tEvents.get(1).getEvent()).getTargetIndex();
  // order of vA task completion can change order of events
  Assert.assertTrue("t1: " + targetIndex1 + " t2: " + targetIndex2,
      (targetIndex1 == 0 && targetIndex2 == 1) || (targetIndex1 == 1 && targetIndex2 == 0));
  vImpl = (VertexImpl) dagImpl.getVertex(vC.getName());
  tImpl = (TaskImpl) vImpl.getTask(1);
  taId = TezTaskAttemptID.getInstance(tImpl.getTaskId(), 0);
  tEvents = vImpl.getTaskAttemptTezEvents(taId, 0, 0, 1000).getEvents();
  Assert.assertEquals(2, tEvents.size()); // 2 from vA
  Assert.assertEquals(vA.getName(), tEvents.get(0).getDestinationInfo().getEdgeVertexName());
  Assert.assertEquals(1, ((CompositeRoutedDataMovementEvent)tEvents.get(0).getEvent()).getSourceIndex());
  Assert.assertEquals(vA.getName(), tEvents.get(1).getDestinationInfo().getEdgeVertexName());
  Assert.assertEquals(1, ((CompositeRoutedDataMovementEvent)tEvents.get(1).getEvent()).getSourceIndex());
  targetIndex1 = ((CompositeRoutedDataMovementEvent)tEvents.get(0).getEvent()).getTargetIndex();
  targetIndex2 = ((CompositeRoutedDataMovementEvent)tEvents.get(1).getEvent()).getTargetIndex();
  // order of vA task completion can change order of events
  Assert.assertTrue("t1: " + targetIndex1 + " t2: " + targetIndex2,
      (targetIndex1 == 0 && targetIndex2 == 1) || (targetIndex1 == 1 && targetIndex2 == 0));
  vImpl = (VertexImpl) dagImpl.getVertex(vD.getName());
  tImpl = (TaskImpl) vImpl.getTask(1);
  taId = TezTaskAttemptID.getInstance(tImpl.getTaskId(), 0);
  tEvents = vImpl.getTaskAttemptTezEvents(taId, 0, 0, 1000).getEvents();
  Assert.assertEquals(1, tEvents.size()); // 1 from vA
  Assert.assertEquals(vA.getName(), tEvents.get(0).getDestinationInfo().getEdgeVertexName());
  Assert.assertEquals(0, ((DataMovementEvent)tEvents.get(0).getEvent()).getTargetIndex());
  Assert.assertEquals(0, ((DataMovementEvent)tEvents.get(0).getEvent()).getSourceIndex());

  tezClient.stop();
}
 

@Test (timeout = 100000)
public void testMixedEdgeRouting() throws Exception {
 TezConfiguration tezconf = new TezConfiguration(defaultConf);
  
  MockTezClient tezClient = new MockTezClient("testMockAM", tezconf, true, null, null, null, null);
  tezClient.start();
  
  MockDAGAppMaster mockApp = tezClient.getLocalClient().getMockApp();
  MockContainerLauncher mockLauncher = mockApp.getContainerLauncher();
  mockLauncher.startScheduling(false);
  mockApp.eventsDelegate = new TestEventsDelegate();
  DAG dag = DAG.create("testMixedEdgeRouting");
  Vertex vA = Vertex.create("A", ProcessorDescriptor.create("Proc.class"), 1);
  Vertex vB = Vertex.create("B", ProcessorDescriptor.create("Proc.class"), 1);
  Vertex vC = Vertex.create("C", ProcessorDescriptor.create("Proc.class"), 1);
  Vertex vD = Vertex.create("D", ProcessorDescriptor.create("Proc.class"), 1);
  Vertex vE = Vertex.create("E", ProcessorDescriptor.create("Proc.class"), 1);
  dag.addVertex(vA)
      .addVertex(vB)
      .addVertex(vC)
      .addVertex(vD)
      .addVertex(vE)
      .addEdge(
          Edge.create(vA, vC, EdgeProperty.create(DataMovementType.SCATTER_GATHER,
              DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
              OutputDescriptor.create("Out"), InputDescriptor.create("In"))))
      .addEdge(
          Edge.create(vB, vC, EdgeProperty.create(DataMovementType.SCATTER_GATHER,
              DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
              OutputDescriptor.create("Out"), InputDescriptor.create("In"))))
      .addEdge(
          Edge.create(vA, vD, EdgeProperty.create(DataMovementType.SCATTER_GATHER,
              DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
              OutputDescriptor.create("Out"), InputDescriptor.create("In"))))
      .addEdge(
          Edge.create(vB, vD, EdgeProperty.create(
              EdgeManagerPluginDescriptor.create(LegacyEdgeTestEdgeManager.class.getName()),
              DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
              OutputDescriptor.create("Out"), InputDescriptor.create("In"))))
      .addEdge(
          Edge.create(vB, vE, EdgeProperty.create(
            EdgeManagerPluginDescriptor.create(LegacyEdgeTestEdgeManager.class.getName()),
            DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
            OutputDescriptor.create("Out"), InputDescriptor.create("In"))));

  DAGClient dagClient = tezClient.submitDAG(dag);
  mockLauncher.waitTillContainersLaunched();
  DAGImpl dagImpl = (DAGImpl) mockApp.getContext().getCurrentDAG();
  mockLauncher.startScheduling(true);
  dagClient.waitForCompletion();
  Assert.assertEquals(DAGStatus.State.SUCCEEDED, dagClient.getDAGStatus(null).getState());
  // vC uses on demand routing and its task does not provide events
  VertexImpl vImpl = (VertexImpl) dagImpl.getVertex(vC.getName());
  TaskImpl tImpl = (TaskImpl) vImpl.getTask(0);
  TezTaskAttemptID taId = TezTaskAttemptID.getInstance(tImpl.getTaskId(), 0);
  Assert.assertEquals(0, tImpl.getTaskAttemptTezEvents(taId, 0, 1000).size());
  // vD is mixed mode and only 1 out of 2 edges does legacy routing with task providing events
  vImpl = (VertexImpl) dagImpl.getVertex(vD.getName());
  tImpl = (TaskImpl) vImpl.getTask(0);
  taId = TezTaskAttemptID.getInstance(tImpl.getTaskId(), 0);
  Assert.assertEquals(1, tImpl.getTaskAttemptTezEvents(taId, 0, 1000).size());
  // vE has single legacy edge and does not use on demand routing and its task provides events
  vImpl = (VertexImpl) dagImpl.getVertex(vE.getName());
  tImpl = (TaskImpl) vImpl.getTask(0);
  taId = TezTaskAttemptID.getInstance(tImpl.getTaskId(), 0);
  Assert.assertEquals(1, tImpl.getTaskAttemptTezEvents(taId, 0, 1000).size());

  tezClient.stop();
}
 

@Test
public void testCountersAggregation() throws Exception {
  TezConfiguration tezconf = new TezConfiguration(defaultConf);
  MockTezClient tezClient = new MockTezClient("testMockAM", tezconf, true, null, null, null,
                                              null, false, false);
  tezClient.start();

  final String vAName = "A";
  final String vBName = "B";
  final String procCounterName = "Proc";
  final String globalCounterName = "Global";
  DAG dag = DAG.create("testCountersAggregation");
  Vertex vA = Vertex.create(vAName, ProcessorDescriptor.create("Proc.class"), 10);
  Vertex vB = Vertex.create(vBName, ProcessorDescriptor.create("Proc.class"), 1);
  dag.addVertex(vA)
      .addVertex(vB)
      .addEdge(
          Edge.create(vA, vB, EdgeProperty.create(DataMovementType.SCATTER_GATHER,
                                                  DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
                                                  OutputDescriptor.create("Out"), InputDescriptor.create("In"))));
  TezCounters temp = new TezCounters();
  temp.findCounter(new String(globalCounterName), new String(globalCounterName)).increment(1);
  ByteArrayOutputStream bos = new ByteArrayOutputStream();
  DataOutput out = new DataOutputStream(bos);
  temp.write(out);
  final byte[] payload = bos.toByteArray();

  MockDAGAppMaster mockApp = tezClient.getLocalClient().getMockApp();
  MockContainerLauncher mockLauncher = mockApp.getContainerLauncher();
  mockLauncher.startScheduling(false);
  mockApp.countersDelegate = new CountersDelegate() {
    int counterValue = 0;
    @Override
    public TezCounters getCounters(TaskSpec taskSpec) {
      String vName = taskSpec.getVertexName();
      TezCounters counters = new TezCounters();
      final DataInputByteBuffer in  = new DataInputByteBuffer();
      in.reset(ByteBuffer.wrap(payload));
      try {
        // this ensures that the serde code path is covered.
        // the internal merges of counters covers the constructor code path.
        counters.readFields(in);
      } catch (IOException e) {
        Assert.fail(e.getMessage());
      }
      counters.findCounter(vName, procCounterName).setValue(++counterValue);
      for (OutputSpec output : taskSpec.getOutputs()) {
        counters.findCounter(vName, output.getDestinationVertexName()).setValue(++counterValue);
      }
      for (InputSpec input : taskSpec.getInputs()) {
        counters.findCounter(vName, input.getSourceVertexName()).setValue(++counterValue);
      }
      return counters;
    }
  };
  mockApp.doSleep = false;
  DAGClient dagClient = tezClient.submitDAG(dag);
  mockLauncher.waitTillContainersLaunched();
  DAGImpl dagImpl = (DAGImpl) mockApp.getContext().getCurrentDAG();
  mockLauncher.startScheduling(true);
  DAGStatus status = dagClient.waitForCompletion();
  Assert.assertEquals(DAGStatus.State.SUCCEEDED, status.getState());
  TezCounters counters = dagImpl.getAllCounters();

  // verify processor counters
  VertexImpl vAImpl = (VertexImpl) dagImpl.getVertex(vAName);
  VertexImpl vBImpl = (VertexImpl) dagImpl.getVertex(vBName);
  TezCounters vACounters = vAImpl.getAllCounters();
  TezCounters vBCounters = vBImpl.getAllCounters();

  Assert.assertEquals(19, ((AggregateTezCounterDelegate)vACounters.findCounter(vAName, procCounterName)).getMax());
  Assert.assertEquals(1, ((AggregateTezCounterDelegate)vACounters.findCounter(vAName, procCounterName)).getMin());
  Assert.assertEquals(20, ((AggregateTezCounterDelegate)vACounters.findCounter(vAName, vBName)).getMax());
  Assert.assertEquals(2, ((AggregateTezCounterDelegate)vACounters.findCounter(vAName, vBName)).getMin());

  Assert.assertEquals(21, ((AggregateTezCounterDelegate)vBCounters.findCounter(vBName, procCounterName)).getMin());
  Assert.assertEquals(21, ((AggregateTezCounterDelegate)vBCounters.findCounter(vBName, procCounterName)).getMax());
  Assert.assertEquals(22, ((AggregateTezCounterDelegate)vBCounters.findCounter(vBName, vAName)).getMin());
  Assert.assertEquals(22, ((AggregateTezCounterDelegate)vBCounters.findCounter(vBName, vAName)).getMax());

  tezClient.stop();
}
 

private DAG createDAG(FileSystem fs, JobID jobId, Configuration[] stageConfs,
    String jobSubmitDir, Credentials ts,
    Map<String, LocalResource> jobLocalResources) throws IOException {

  String jobName = stageConfs[0].get(MRJobConfig.JOB_NAME,
      YarnConfiguration.DEFAULT_APPLICATION_NAME);
  DAG dag = DAG.create(jobName);

  LOG.info("Number of stages: " + stageConfs.length);

  List<TaskLocationHint> mapInputLocations =
      getMapLocationHintsFromInputSplits(
          jobId, fs, stageConfs[0], jobSubmitDir);
  List<TaskLocationHint> reduceInputLocations = null;

  Vertex[] vertices = new Vertex[stageConfs.length];
  for (int i = 0; i < stageConfs.length; i++) {
    vertices[i] = createVertexForStage(stageConfs[i], jobLocalResources,
        i == 0 ? mapInputLocations : reduceInputLocations, i,
        stageConfs.length);
  }

  for (int i = 0; i < vertices.length; i++) {
    dag.addVertex(vertices[i]);
    if (i > 0) {
      // Set edge conf based on Input conf (compression etc properties for MapReduce are
      // w.r.t Outputs - MAP_OUTPUT_COMPRESS for example)
      Map<String, String> partitionerConf = null;
      if (stageConfs[i-1] != null) {
        partitionerConf = Maps.newHashMap();
        for (Map.Entry<String, String> entry : stageConfs[i - 1]) {
          partitionerConf.put(entry.getKey(), entry.getValue());
        }
      }
      OrderedPartitionedKVEdgeConfig edgeConf =
          OrderedPartitionedKVEdgeConfig.newBuilder(stageConfs[i - 1].get(
                  TezRuntimeConfiguration.TEZ_RUNTIME_KEY_CLASS),
              stageConfs[i - 1].get(TezRuntimeConfiguration.TEZ_RUNTIME_VALUE_CLASS),
              MRPartitioner.class.getName(), partitionerConf)
              .setFromConfigurationUnfiltered(stageConfs[i-1])
              .configureInput().useLegacyInput().done()
              .build();
      Edge edge = Edge.create(vertices[i - 1], vertices[i], edgeConf.createDefaultEdgeProperty());
      dag.addEdge(edge);
    }

  }
  return dag;
}
 

@Test (timeout = 10000)
public void testBasicCounters() throws Exception {
  TezConfiguration tezconf = new TezConfiguration(defaultConf);
  MockTezClient tezClient = new MockTezClient("testMockAM", tezconf, true, null, null, null,
      null, false, false);
  tezClient.start();

  final String vAName = "A";
  final String vBName = "B";
  final String procCounterName = "Proc";
  final String globalCounterName = "Global";
  DAG dag = DAG.create("testBasicCounters");
  Vertex vA = Vertex.create(vAName, ProcessorDescriptor.create("Proc.class"), 10);
  Vertex vB = Vertex.create(vBName, ProcessorDescriptor.create("Proc.class"), 1);
  dag.addVertex(vA)
      .addVertex(vB)
      .addEdge(
          Edge.create(vA, vB, EdgeProperty.create(DataMovementType.SCATTER_GATHER,
              DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
              OutputDescriptor.create("Out"), InputDescriptor.create("In"))));
  TezCounters temp = new TezCounters();
  temp.findCounter(new String(globalCounterName), new String(globalCounterName)).increment(1);
  ByteArrayOutputStream bos = new ByteArrayOutputStream();
  DataOutput out = new DataOutputStream(bos);
  temp.write(out);
  final byte[] payload = bos.toByteArray();

  MockDAGAppMaster mockApp = tezClient.getLocalClient().getMockApp();
  MockContainerLauncher mockLauncher = mockApp.getContainerLauncher();
  mockLauncher.startScheduling(false);
  mockApp.countersDelegate = new CountersDelegate() {
    @Override
    public TezCounters getCounters(TaskSpec taskSpec) {
      String vName = taskSpec.getVertexName();
      TezCounters counters = new TezCounters();
      final DataInputByteBuffer in  = new DataInputByteBuffer();
      in.reset(ByteBuffer.wrap(payload));
      try {
        // this ensures that the serde code path is covered.
        // the internal merges of counters covers the constructor code path.
        counters.readFields(in);
      } catch (IOException e) {
        Assert.fail(e.getMessage());
      }
      counters.findCounter(vName, procCounterName).increment(1);
      for (OutputSpec output : taskSpec.getOutputs()) {
        counters.findCounter(vName, output.getDestinationVertexName()).increment(1);
      }
      for (InputSpec input : taskSpec.getInputs()) {
        counters.findCounter(vName, input.getSourceVertexName()).increment(1);
      }
      return counters;
    }
  };
  mockApp.doSleep = false;
  DAGClient dagClient = tezClient.submitDAG(dag);
  mockLauncher.waitTillContainersLaunched();
  DAGImpl dagImpl = (DAGImpl) mockApp.getContext().getCurrentDAG();
  mockLauncher.startScheduling(true);
  DAGStatus status = dagClient.waitForCompletion();
  Assert.assertEquals(DAGStatus.State.SUCCEEDED, status.getState());
  TezCounters counters = dagImpl.getAllCounters();

  String osName = System.getProperty("os.name").toLowerCase(Locale.ENGLISH);
  if (SystemUtils.IS_OS_LINUX) {
    Assert.assertTrue(counters.findCounter(DAGCounter.AM_CPU_MILLISECONDS).getValue() > 0);
  }

  // verify processor counters
  Assert.assertEquals(10, counters.findCounter(vAName, procCounterName).getValue());
  Assert.assertEquals(1, counters.findCounter(vBName, procCounterName).getValue());
  // verify edge counters
  Assert.assertEquals(10, counters.findCounter(vAName, vBName).getValue());
  Assert.assertEquals(1, counters.findCounter(vBName, vAName).getValue());
  // verify global counters
  Assert.assertEquals(11, counters.findCounter(globalCounterName, globalCounterName).getValue());
  VertexImpl vAImpl = (VertexImpl) dagImpl.getVertex(vAName);
  VertexImpl vBImpl = (VertexImpl) dagImpl.getVertex(vBName);
  TezCounters vACounters = vAImpl.getAllCounters();
  TezCounters vBCounters = vBImpl.getAllCounters();
  String vACounterName = vACounters.findCounter(globalCounterName, globalCounterName).getName();
  String vBCounterName = vBCounters.findCounter(globalCounterName, globalCounterName).getName();
  if (vACounterName != vBCounterName) {
    Assert.fail("String counter name objects dont match despite interning.");
  }
  CounterGroup vaGroup = vACounters.getGroup(globalCounterName);
  String vaGrouName = vaGroup.getName();
  CounterGroup vBGroup = vBCounters.getGroup(globalCounterName);
  String vBGrouName = vBGroup.getName();
  if (vaGrouName != vBGrouName) {
    Assert.fail("String group name objects dont match despite interning.");
  }
  
  tezClient.stop();
}