org.apache.tez.common.counters.TezCounters Java Examples

/**
 * Fine-grained recovery task-level, In a vertex (v1), task 0 is done task 1
 * is not started. History flush happens. AM dies. Once AM is recovered, task 0 is
 * not re-run. Task 1 is re-run. (ONE_TO_ONE)
 *
 * @throws Exception
 */
@Test(timeout = 120000)
public void testVertexPartialFinished_One2One() throws Exception {
  DAG dag =
      createDAG("VertexPartialFinished_One2One", ControlledInputReadyVertexManager.class,
          DataMovementType.ONE_TO_ONE, true);
  TezCounters counters = runDAGAndVerify(dag, DAGStatus.State.SUCCEEDED);
  assertEquals(4, counters.findCounter(DAGCounter.NUM_SUCCEEDED_TASKS).getValue());
  assertEquals(2, counters.findCounter(TestCounter.Counter_1).getValue());

  List<HistoryEvent> historyEvents1 = readRecoveryLog(1);
  List<HistoryEvent> historyEvents2 = readRecoveryLog(2);
  printHistoryEvents(historyEvents1, 1);
  printHistoryEvents(historyEvents1, 2);
  // task_0 of v1 is finished in attempt 1, task_1 of v1 is not finished in
  // attempt 1
  assertEquals(1, findTaskAttemptFinishedEvent(historyEvents1, 0, 0).size());
  assertEquals(0, findTaskAttemptFinishedEvent(historyEvents1, 0, 1).size());

  // task_0 of v1 is finished in attempt 1 and not rerun, task_1 of v1 is
  // finished in attempt 2
  assertEquals(1, findTaskAttemptFinishedEvent(historyEvents2, 0, 0).size());
  assertEquals(1, findTaskAttemptFinishedEvent(historyEvents2, 0, 1).size());

}
 

public static Counters fromTez(TezCounters tezCounters) {
  if (tezCounters == null) {
    return null;
  }
  Counters counters = new Counters();
  for (CounterGroup xGrp : tezCounters) {
    counters.addGroup(xGrp.getName(), xGrp.getDisplayName());
    for (TezCounter xCounter : xGrp) {
      Counter counter =
          counters.findCounter(xGrp.getName(), xCounter.getName());
      counter.setValue(xCounter.getValue());

    }
  }
  return counters;
}
 

void logJobHistoryFinishedEvent(TezCounters counters) throws IOException {
  if (recoveryData == null
      || recoveryData.getDAGFinishedEvent() == null) {
    Map<String, Integer> taskStats = constructTaskStats(getDAGProgress());
    if (finishTime < startTime) {
      LOG.warn("DAG finish time is smaller than start time. "
          + "startTime=" + startTime
          + ", finishTime=" + finishTime
      );
    }
    DAGFinishedEvent finishEvt = new DAGFinishedEvent(dagId, startTime,
        finishTime, DAGState.SUCCEEDED, "", counters,
        this.userName, this.dagName, taskStats, this.appContext.getApplicationAttemptId(),
        this.jobPlan);
    this.appContext.getHistoryHandler().handleCriticalEvent(
        new DAGHistoryEvent(dagId, finishEvt));
  }
}
 

@Override
public int hashCode() {
  final int prime = 46021;
  int result = 1;
  result = prime +
      getState().hashCode();

  List<String> diagnostics = getDiagnostics();
  Progress vProgress = getProgress();
  TezCounters counters = getVertexCounters();

  result = prime * result +
      ((diagnostics == null)? 0 : diagnostics.hashCode());
  result = prime * result +
      ((vProgress == null)? 0 : vProgress.hashCode());
  result = prime * result +
      ((counters == null)? 0 : counters.hashCode());

  return result;
}
 

private float getAvgCounter(Collection<TaskAttemptInfo> taskAttemptInfos, String
    counterGroupName, String counterName) {
  long total = 0;
  int taskCount = 0;
  for (TaskAttemptInfo attemptInfo : taskAttemptInfos) {
    TezCounters tezCounters = attemptInfo.getTezCounters();
    TezCounter counter = tezCounters.findCounter(counterGroupName, counterName);
    if (counter != null) {
      total += counter.getValue();
      taskCount++;
    } else {
      LOG.info("Could not find counterGroupName=" + counterGroupName + ", counter=" +
          counterName + " in " + attemptInfo);
    }
  }
  return (taskCount > 0) ? (total * 1.0f / taskCount) : 0;
}
 

@Override
public TezCounters getAllCounters() {

  readLock.lock();

  try {
    DAGState state = getInternalState();
    if (state == DAGState.ERROR || state == DAGState.FAILED
        || state == DAGState.KILLED || state == DAGState.SUCCEEDED) {
      this.mayBeConstructFinalFullCounters();
      return fullCounters;
    }

    // dag not yet finished. update cpu time counters
    updateCpuCounters();
    TezCounters counters = new TezCounters();
    counters.incrAllCounters(dagCounters);
    return aggrTaskCounters(counters, vertices.values());

  } finally {
    readLock.unlock();
  }
}
 

public static Map<String,Object> convertCountersToATSMap(TezCounters counters) {
  Map<String,Object> object = new LinkedHashMap<String, Object>();
  if (counters == null) {
      return object;
    }
  ArrayList<Object> counterGroupsList = new ArrayList<Object>();
  for (CounterGroup group : counters) {
      Map<String,Object> counterGroupMap = new LinkedHashMap<String, Object>();
      counterGroupMap.put(ATSConstants.COUNTER_GROUP_NAME, group.getName());
      counterGroupMap.put(ATSConstants.COUNTER_GROUP_DISPLAY_NAME,
              group.getDisplayName());
      ArrayList<Object> counterList = new ArrayList<Object>();
      for (TezCounter counter : group) {
          Map<String,Object> counterMap = new LinkedHashMap<String, Object>();
          counterMap.put(ATSConstants.COUNTER_NAME, counter.getName());
          counterMap.put(ATSConstants.COUNTER_DISPLAY_NAME,
                  counter.getDisplayName());
          counterMap.put(ATSConstants.COUNTER_VALUE, counter.getValue());
          counterList.add(counterMap);
        }
      putInto(counterGroupMap, ATSConstants.COUNTERS, counterList);
      counterGroupsList.add(counterGroupMap);
    }
  putInto(object, ATSConstants.COUNTER_GROUPS, counterGroupsList);
  return object;
}
 

private InputContext createTezInputContext(Configuration payloadConf, Configuration baseConf) throws Exception {
  MRInputUserPayloadProto.Builder builder = MRInputUserPayloadProto.newBuilder();
  builder.setGroupingEnabled(false);
  builder.setConfigurationBytes(TezUtils.createByteStringFromConf(payloadConf));
  byte[] payload = builder.build().toByteArray();

  ApplicationId applicationId = ApplicationId.newInstance(10000, 1);
  TezCounters counters = new TezCounters();

  InputContext inputContext = mock(InputContext.class);
  doReturn(applicationId).when(inputContext).getApplicationId();
  doReturn(counters).when(inputContext).getCounters();
  doReturn(1).when(inputContext).getDAGAttemptNumber();
  doReturn("dagName").when(inputContext).getDAGName();
  doReturn(1).when(inputContext).getInputIndex();
  doReturn("srcVertexName").when(inputContext).getSourceVertexName();
  doReturn(1).when(inputContext).getTaskAttemptNumber();
  doReturn(1).when(inputContext).getTaskIndex();
  doReturn(1).when(inputContext).getTaskVertexIndex();
  doReturn(UUID.randomUUID().toString()).when(inputContext).getUniqueIdentifier();
  doReturn("taskVertexName").when(inputContext).getTaskVertexName();
  doReturn(UserPayload.create(ByteBuffer.wrap(payload))).when(inputContext).getUserPayload();
  doReturn(baseConf).when(inputContext).getContainerConfiguration();
  return inputContext;
}
 

@Test(timeout = 5000)
public void testTezClientCounterLimits() throws YarnException, IOException, ServiceException {
  Limits.reset();
  int defaultCounterLimit = TezConfiguration.TEZ_COUNTERS_MAX_DEFAULT;

  int newCounterLimit = defaultCounterLimit + 500;

  TezConfiguration conf = new TezConfiguration();
  conf.setInt(TezConfiguration.TEZ_COUNTERS_MAX, newCounterLimit);

  configureAndCreateTezClient(conf);

  TezCounters counters = new TezCounters();
  for (int i = 0 ; i < newCounterLimit ; i++) {
    counters.findCounter("GroupName", "TestCounter" + i).setValue(i);
  }

  try {
    counters.findCounter("GroupName", "TestCounterFail").setValue(1);
    fail("Expecting a LimitExceedException - too many counters");
  } catch (LimitExceededException e) {
  }
}
 

public VertexFinishedEvent(TezVertexID vertexId, String vertexName, int numTasks, long initRequestedTime,
                           long initedTime, long startRequestedTime, long startedTime,
                           long finishTime, VertexState state, String diagnostics,
                           TezCounters counters, VertexStats vertexStats,
                           Map<String, Integer> vertexTaskStats,
                           ServicePluginInfo servicePluginInfo) {
  this.vertexName = vertexName;
  this.vertexID = vertexId;
  this.numTasks = numTasks;
  this.initRequestedTime = initRequestedTime;
  this.initedTime = initedTime;
  this.startRequestedTime = startRequestedTime;
  this.startTime = startedTime;
  this.finishTime = finishTime;
  this.state = state;
  this.diagnostics = diagnostics;
  this.tezCounters = counters;
  this.vertexStats = vertexStats;
  this.vertexTaskStats = vertexTaskStats;
  this.servicePluginInfo = servicePluginInfo;
}
 

@Override
public TezCounters getAllCounters() {

  readLock.lock();

  try {
    if (inTerminalState()) {
      this.mayBeConstructFinalFullCounters();
      return fullCounters;
    }

    TezCounters counters = new TezCounters();
    counters.aggrAllCounters(this.counters);
    return aggrTaskCounters(counters, tasks.values());

  } finally {
    readLock.unlock();
  }
}
 

public TezProcessorContextImpl(Configuration conf, String[] workDirs, int appAttemptNumber,
    TezUmbilical tezUmbilical, String dagName, String vertexName,
    TezTaskAttemptID taskAttemptID, TezCounters counters,
    @Nullable byte[] userPayload, RuntimeTask runtimeTask,
    Map<String, ByteBuffer> serviceConsumerMetadata,
    Map<String, String> auxServiceEnv, MemoryDistributor memDist,
    ProcessorDescriptor processorDescriptor, InputReadyTracker inputReadyTracker) {
  super(conf, workDirs, appAttemptNumber, dagName, vertexName, taskAttemptID,
      counters, runtimeTask, tezUmbilical, serviceConsumerMetadata,
      auxServiceEnv, memDist, processorDescriptor);
  checkNotNull(inputReadyTracker, "inputReadyTracker is null");
  this.userPayload = DagTypeConverters.convertToTezUserPayload(userPayload);
  this.sourceInfo = new EventMetaData(EventProducerConsumerType.PROCESSOR,
      taskVertexName, "", taskAttemptID);
  this.inputReadyTracker = inputReadyTracker;
}
 

@Private
public TezInputContextImpl(Configuration conf, String[] workDirs, int appAttemptNumber,
    TezUmbilical tezUmbilical, String dagName, String taskVertexName,
    String sourceVertexName, TezTaskAttemptID taskAttemptID,
    TezCounters counters, int inputIndex, @Nullable byte[] userPayload,
    RuntimeTask runtimeTask, Map<String, ByteBuffer> serviceConsumerMetadata,
    Map<String, String> auxServiceEnv, MemoryDistributor memDist,
    InputDescriptor inputDescriptor,  Input input, InputReadyTracker inputReadyTracker) {
  super(conf, workDirs, appAttemptNumber, dagName, taskVertexName, taskAttemptID,
      wrapCounters(counters, taskVertexName, sourceVertexName, conf),
      runtimeTask, tezUmbilical, serviceConsumerMetadata,
      auxServiceEnv, memDist, inputDescriptor);
  checkNotNull(inputIndex, "inputIndex is null");
  checkNotNull(sourceVertexName, "sourceVertexName is null");
  checkNotNull(input, "input is null");
  checkNotNull(inputReadyTracker, "inputReadyTracker is null");
  this.userPayload = DagTypeConverters.convertToTezUserPayload(userPayload);
  this.inputIndex = inputIndex;
  this.sourceVertexName = sourceVertexName;
  this.sourceInfo = new EventMetaData(
      EventProducerConsumerType.INPUT, taskVertexName, sourceVertexName,
      taskAttemptID);
  this.input = input;
  this.inputReadyTracker = inputReadyTracker;
}
 

static OutputContext createOutputContext() throws IOException {
  OutputContext outputContext = mock(OutputContext.class);
  Configuration conf = new TezConfiguration();
  UserPayload payLoad = TezUtils.createUserPayloadFromConf(conf);
  String[] workingDirs = new String[]{"workDir1"};
  OutputStatisticsReporter statsReporter = mock(OutputStatisticsReporter.class);
  TezCounters counters = new TezCounters();

  doReturn("destinationVertex").when(outputContext).getDestinationVertexName();
  doReturn(payLoad).when(outputContext).getUserPayload();
  doReturn(workingDirs).when(outputContext).getWorkDirs();
  doReturn(200 * 1024 * 1024l).when(outputContext).getTotalMemoryAvailableToTask();
  doReturn(counters).when(outputContext).getCounters();
  doReturn(statsReporter).when(outputContext).getStatisticsReporter();
  doReturn(new Configuration(false)).when(outputContext).getContainerConfiguration();
  return outputContext;
}
 

@Override
public int hashCode() {
  // Source explicitly exclude from hashCode
  final int prime = 44017;
  int result = 1;
  result = prime +
      getState().hashCode();

  List<String> diagnostics = getDiagnostics();
  Progress dagProgress = getDAGProgress();
  Map<String, Progress> vProgress = getVertexProgress();
  TezCounters counters = getDAGCounters();

  result = prime * result +
      ((diagnostics == null)? 0 : diagnostics.hashCode());
  result = prime * result +
      ((dagProgress == null)? 0 : dagProgress.hashCode());
  result = prime * result +
      ((vProgress == null)? 0 : vProgress.hashCode());
  result = prime * result +
      ((counters == null)? 0 : counters.hashCode());

  return result;
}
 

/**
 * Fine-grained recovery task-level, In a vertex (v1), task 0 is done task 1
 * is also done. History flush happens. AM dies. Once AM is recovered, task 0
 * and Task 1 is not re-run. (Broadcast)
 *
 * @throws Exception
 */
@Test(timeout = 120000)
public void testVertexCompletelyFinished_Broadcast() throws Exception {
  DAG dag =
      createDAG("VertexCompletelyFinished_Broadcast", ControlledImmediateStartVertexManager.class,
          DataMovementType.BROADCAST, false);
  TezCounters counters = runDAGAndVerify(dag, DAGStatus.State.SUCCEEDED);

  assertEquals(4, counters.findCounter(DAGCounter.NUM_SUCCEEDED_TASKS).getValue());
  assertEquals(2, counters.findCounter(TestCounter.Counter_1).getValue());

  List<HistoryEvent> historyEvents1 = readRecoveryLog(1);
  List<HistoryEvent> historyEvents2 = readRecoveryLog(2);
  printHistoryEvents(historyEvents1, 1);
  printHistoryEvents(historyEvents1, 2);
  // task_0 of v1 is finished in attempt 1, task_1 of v1 is not finished in
  // attempt 1
  assertEquals(1, findTaskAttemptFinishedEvent(historyEvents1, 0, 0).size());
  assertEquals(1, findTaskAttemptFinishedEvent(historyEvents1, 0, 1).size());

  // task_0 of v1 is finished in attempt 1 and not rerun, task_1 of v1 is
  // finished in attempt 2
  assertEquals(1, findTaskAttemptFinishedEvent(historyEvents2, 0, 0).size());
  assertEquals(1, findTaskAttemptFinishedEvent(historyEvents2, 0, 1).size());
}
 

public VertexFinishedEvent(TezVertexID vertexId,
    String vertexName, long initRequestedTime, long initedTime,
    long startRequestedTime, long startedTime, long finishTime,
    VertexState state, String diagnostics, TezCounters counters,
    VertexStats vertexStats) {
  this.vertexName = vertexName;
  this.vertexID = vertexId;
  this.initRequestedTime = initRequestedTime;
  this.initedTime = initedTime;
  this.startRequestedTime = startRequestedTime;
  this.startTime = startedTime;
  this.finishTime = finishTime;
  this.state = state;
  this.diagnostics = diagnostics;
  this.tezCounters = counters;
  this.vertexStats = vertexStats;
}
 

private InputContext createTezInputContext() throws IOException {
  ApplicationId applicationId = ApplicationId.newInstance(1, 1);
  InputContext inputContext = mock(InputContext.class);
  doReturn(applicationId).when(inputContext).getApplicationId();
  doReturn("sourceVertex").when(inputContext).getSourceVertexName();
  when(inputContext.getCounters()).thenReturn(new TezCounters());
  ExecutionContext executionContext = new ExecutionContextImpl("localhost");
  doReturn(executionContext).when(inputContext).getExecutionContext();
  ByteBuffer shuffleBuffer = ByteBuffer.allocate(4).putInt(0, 4);
  doReturn(shuffleBuffer).when(inputContext).getServiceProviderMetaData(anyString());
  Token<JobTokenIdentifier>
      sessionToken = new Token<JobTokenIdentifier>(new JobTokenIdentifier(new Text("text")),
      new JobTokenSecretManager());
  ByteBuffer tokenBuffer = TezCommonUtils.serializeServiceData(sessionToken);
  doReturn(tokenBuffer).when(inputContext).getServiceConsumerMetaData(anyString());
  when(inputContext.createTezFrameworkExecutorService(anyInt(), anyString())).thenAnswer(
      new Answer<ExecutorService>() {
        @Override
        public ExecutorService answer(InvocationOnMock invocation) throws Throwable {
          return sharedExecutor.createExecutorService(
              invocation.getArgumentAt(0, Integer.class),
              invocation.getArgumentAt(1, String.class));
        }
      });
  return inputContext;
}
 

public TaskStatusUpdateEvent(TezCounters tezCounters, float progress, TaskStatistics statistics, 
    boolean progressNotified) {
  this.tezCounters = tezCounters;
  this.progress = progress;
  this.statistics = statistics;
  this.progressNotified = progressNotified;
}
 

@Private
@VisibleForTesting
void setCounters(TezCounters counters) {
  try {
    writeLock.lock();
    this.counters = counters;
  } finally {
    writeLock.unlock();
  }
}
 

private OutputContext createMockOutputContext(UserPayload payload, Configuration baseConf) {
  OutputContext outputContext = mock(OutputContext.class);
  ApplicationId appId = ApplicationId.newInstance(System.currentTimeMillis(), 1);
  when(outputContext.getUserPayload()).thenReturn(payload);
  when(outputContext.getApplicationId()).thenReturn(appId);
  when(outputContext.getTaskVertexIndex()).thenReturn(1);
  when(outputContext.getTaskAttemptNumber()).thenReturn(1);
  when(outputContext.getCounters()).thenReturn(new TezCounters());
  when(outputContext.getContainerConfiguration()).thenReturn(baseConf);
  return outputContext;
}
 

@Test(timeout = 5000)
public void testGetVertexStatusSimple() throws JSONException, TezException, IOException {
  DAGClientTimelineImpl
      httpClient = new DAGClientTimelineImpl(mock(ApplicationId.class), "EXAMPLE_DAG_ID",
      new TezConfiguration(), null, 0);
  DAGClientTimelineImpl spyClient = spy(httpClient);
  spyClient.baseUri = "http://yarn.ats.webapp/ws/v1/timeline";
  final String expectedVertexUrl = "http://yarn.ats.webapp/ws/v1/timeline/TEZ_VERTEX_ID" +
      "?primaryFilter=TEZ_DAG_ID:EXAMPLE_DAG_ID&secondaryFilter=vertexName:vertex1name&" +
      "fields=primaryfilters,otherinfo";

  Set<StatusGetOpts> statusOptions = new HashSet<StatusGetOpts>(1);
  statusOptions.add(StatusGetOpts.GET_COUNTERS);

  final String jsonData = "{entities:[ {otherinfo:{numFailedTasks:1,numSucceededTasks:2," +
      "status:'SUCCEEDED', vertexName:'vertex1name', numTasks:4, numKilledTasks: 3, " +
      "numCompletedTasks: 4, diagnostics: 'diagnostics1', " +
      "counters: { counterGroups: [ " +
      "      { counterGroupName: 'CG1', counterGroupDisplayName: 'CGD1', counters: [" +
      "        {counterName:'C1', counterDisplayName: 'CD1', counterValue: 1 }," +
      "        {counterName:'C2', counterDisplayName: 'CD2', counterValue: 2 }" +
      "      ]}" +
      "    ]}" +
      "}}]}";

  doReturn(new JSONObject(jsonData)).when(spyClient).getJsonRootEntity(expectedVertexUrl);

  VertexStatus vertexStatus = spyClient.getVertexStatus("vertex1name", statusOptions);
  Assert.assertEquals("status check", VertexStatus.State.SUCCEEDED, vertexStatus.getState());
  Assert.assertEquals("diagnostics", "diagnostics1", vertexStatus.getDiagnostics().get(0));
  final Progress progress = vertexStatus.getProgress();
  final TezCounters vertexCounters = vertexStatus.getVertexCounters();
  Assert.assertEquals("failed task count", 1, progress.getFailedTaskCount());
  Assert.assertEquals("suceeded task count", 2, progress.getSucceededTaskCount());
  Assert.assertEquals("killed task count", 3, progress.getKilledTaskCount());
  Assert.assertEquals("total task count", 4, progress.getTotalTaskCount());
  Assert.assertEquals("Counters Size", 2, vertexCounters.countCounters());
  Assert.assertEquals("Counter Value", 1,
      vertexCounters.getGroup("CG1").findCounter("C1").getValue());
}
 

private OutputContext createMockOutputContext(TezCounters counters, ApplicationId appId,
    String uniqueId, String auxiliaryService) {
  OutputContext outputContext = mock(OutputContext.class);
  doReturn(counters).when(outputContext).getCounters();
  doReturn(appId).when(outputContext).getApplicationId();
  doReturn(1).when(outputContext).getDAGAttemptNumber();
  doReturn("dagName").when(outputContext).getDAGName();
  doReturn("destinationVertexName").when(outputContext).getDestinationVertexName();
  doReturn(1).when(outputContext).getOutputIndex();
  doReturn(1).when(outputContext).getTaskAttemptNumber();
  doReturn(1).when(outputContext).getTaskIndex();
  doReturn(1).when(outputContext).getTaskVertexIndex();
  doReturn("vertexName").when(outputContext).getTaskVertexName();
  doReturn(uniqueId).when(outputContext).getUniqueIdentifier();

  doAnswer(new Answer<ByteBuffer>() {
    @Override
    public ByteBuffer answer(InvocationOnMock invocation) throws Throwable {
      ByteBuffer portBuffer = ByteBuffer.allocate(4);
      portBuffer.mark();
      portBuffer.putInt(SHUFFLE_PORT);
      portBuffer.reset();
      return portBuffer;
    }
  }).when(outputContext).getServiceProviderMetaData(eq(auxiliaryService));

  Path outDirBase = new Path(TEST_ROOT_DIR, "outDir_" + uniqueId);
  String[] outDirs = new String[] { outDirBase.toString() };
  doReturn(outDirs).when(outputContext).getWorkDirs();
  return outputContext;
}
 

private Map<String, Object> getVertexInfoMap(Vertex vertex,
                                             Map<String, Set<String>> counterNames) {
  Map<String, Object> vertexInfo = new HashMap<String, Object>();
  vertexInfo.put("id", vertex.getVertexId().toString());
  vertexInfo.put("status", vertex.getState().toString());
  vertexInfo.put("progress", Float.toString(vertex.getCompletedTaskProgress()));

  vertexInfo.put("initTime", Long.toString(vertex.getInitTime()));
  vertexInfo.put("startTime", Long.toString(vertex.getStartTime()));
  vertexInfo.put("finishTime", Long.toString(vertex.getFinishTime()));
  vertexInfo.put("firstTaskStartTime", Long.toString(vertex.getFirstTaskStartTime()));
  vertexInfo.put("lastTaskFinishTime", Long.toString(vertex.getLastTaskFinishTime()));

  ProgressBuilder vertexProgress = vertex.getVertexProgress();
  vertexInfo.put("totalTasks", Integer.toString(vertexProgress.getTotalTaskCount()));
  vertexInfo.put("runningTasks", Integer.toString(vertexProgress.getRunningTaskCount()));
  vertexInfo.put("succeededTasks", Integer.toString(vertexProgress.getSucceededTaskCount()));

  vertexInfo.put("failedTaskAttempts",
      Integer.toString(vertexProgress.getFailedTaskAttemptCount()));
  vertexInfo.put("killedTaskAttempts",
      Integer.toString(vertexProgress.getKilledTaskAttemptCount()));

  try {
    if (counterNames != null && !counterNames.isEmpty()) {
      TezCounters counters = vertex.getCachedCounters();
      Map<String, Map<String, Long>> counterMap = constructCounterMapInfo(counters, counterNames);
      if (counterMap != null && !counterMap.isEmpty()) {
        vertexInfo.put("counters", counterMap);
      }
    }
  } catch (LimitExceededException e) {
    // Ignore
    // TODO: add an error message instead for counter key
  }

  return vertexInfo;
}
 

private static TezCounters wrapCounters(TezCounters tezCounters, String taskVertexName,
    String edgeVertexName, Configuration conf) {
  if (conf.getBoolean(TezConfiguration.TEZ_TASK_GENERATE_COUNTERS_PER_IO,
      TezConfiguration.TEZ_TASK_GENERATE_COUNTERS_PER_IO_DEFAULT)) {
    return new TezCountersDelegate(tezCounters, taskVertexName, edgeVertexName, "OUTPUT");
  } else {
    return tezCounters;
  }
}
 

private Vertex createMockVertex(String vertexIDStr, VertexState status, float progress,
                                int taskCounts) {
  ProgressBuilder pb = new ProgressBuilder();
  pb.setTotalTaskCount(taskCounts);
  pb.setSucceededTaskCount(taskCounts * 2);
  pb.setFailedTaskAttemptCount(taskCounts * 3);
  pb.setKilledTaskAttemptCount(taskCounts * 4);
  pb.setRunningTaskCount(taskCounts * 5);

  Vertex mockVertex = mock(Vertex.class);
  doReturn(TezVertexID.fromString(vertexIDStr)).when(mockVertex).getVertexId();
  doReturn(status).when(mockVertex).getState();
  doReturn(progress).when(mockVertex).getProgress();
  doReturn(pb).when(mockVertex).getVertexProgress();
  doReturn(1L).when(mockVertex).getInitTime();
  doReturn(1L).when(mockVertex).getStartTime();
  doReturn(2L).when(mockVertex).getFinishTime();
  doReturn(1L).when(mockVertex).getFirstTaskStartTime();
  doReturn(2L).when(mockVertex).getLastTaskFinishTime();

  TezCounters counters = new TezCounters();
  counters.addGroup("g1", "g1");
  counters.addGroup("g2", "g2");
  counters.addGroup("g3", "g3");
  counters.addGroup("g4", "g4");
  counters.findCounter("g1", "g1_c1").setValue(100);
  counters.findCounter("g1", "g1_c2").setValue(100);
  counters.findCounter("g2", "g2_c3").setValue(100);
  counters.findCounter("g2", "g2_c4").setValue(100);
  counters.findCounter("g3", "g3_c5").setValue(100);
  counters.findCounter("g3", "g3_c6").setValue(100);

  doReturn(counters).when(mockVertex).getAllCounters();
  doReturn(counters).when(mockVertex).getCachedCounters();

  return mockVertex;
}
 

@Test
public void testMergeConf() throws Exception {
  JobConf payloadConf = new JobConf();
  payloadConf.set("local-key", "local-value");
  DataSinkDescriptor dataSink = MultiMROutput.createConfigBuilder(
      payloadConf, SequenceFileOutputFormat.class, "/output", false).build();

  Configuration baseConf = new Configuration(false);
  baseConf.set("base-key", "base-value");

  OutputContext outputContext = mock(OutputContext.class);
  ApplicationId appId = ApplicationId.newInstance(System.currentTimeMillis(), 1);
  when(outputContext.getUserPayload()).thenReturn(dataSink.getOutputDescriptor().getUserPayload());
  when(outputContext.getApplicationId()).thenReturn(appId);
  when(outputContext.getTaskVertexIndex()).thenReturn(1);
  when(outputContext.getTaskAttemptNumber()).thenReturn(1);
  when(outputContext.getCounters()).thenReturn(new TezCounters());
  when(outputContext.getStatisticsReporter()).thenReturn(mock(OutputStatisticsReporter.class));
  when(outputContext.getContainerConfiguration()).thenReturn(baseConf);

  MultiMROutput output = new MultiMROutput(outputContext, 2);
  output.initialize();

  Configuration mergedConf = output.jobConf;
  assertEquals("base-value", mergedConf.get("base-key"));
  assertEquals("local-value", mergedConf.get("local-key"));
}
 

public static TezCounters incrTaskCounters(
    TezCounters counters, Collection<Task> tasks) {
  for (Task task : tasks) {
    counters.incrAllCounters(task.getCounters());
  }
  return counters;
}
 

public static TezCounters incrTaskCounters(
    TezCounters counters, Collection<Vertex> vertices) {
  for (Vertex vertex : vertices) {
    counters.incrAllCounters(vertex.getAllCounters());
  }
  return counters;
}
 

private static TezCounters wrapCounters(LogicalIOProcessorRuntimeTask runtimeTask, String taskVertexName,
    String edgeVertexName, Configuration conf) {
  TezCounters tezCounters = runtimeTask.addAndGetTezCounter(edgeVertexName);
  if (conf.getBoolean(TezConfiguration.TEZ_TASK_GENERATE_COUNTERS_PER_IO,
      TezConfiguration.TEZ_TASK_GENERATE_COUNTERS_PER_IO_DEFAULT)) {
    return new TezCountersDelegate(tezCounters, taskVertexName, edgeVertexName, "OUTPUT");
  } else {
    return tezCounters;
  }
}