TEZ-2119: Counter for launched containers

abstractdog · abstractdog · commit bdccfbe459d7 · 2023-07-04T12:42:29.000+02:00
diff --git a/tez-api/src/main/java/org/apache/tez/common/counters/DAGCounter.java b/tez-api/src/main/java/org/apache/tez/common/counters/DAGCounter.java
@@ -41,5 +41,51 @@ public enum DAGCounter {
   AM_CPU_MILLISECONDS,
   /** Wall clock time taken by all the tasks. */
   WALL_CLOCK_MILLIS,
-  AM_GC_TIME_MILLIS
+  AM_GC_TIME_MILLIS,
+
+  /*
+   * Type: # of containers
+   * Both allocated and launched containers before DAG start.
+   * This is incremented only once when the DAG starts and it's calculated
+   * by querying all the held containers from TaskSchedulers.
+   */
+  INITIAL_HELD_CONTAINERS,
+
+  /*
+   * Type: # of containers
+   * All containers that have been seen/used in this DAG by task allocation.
+   * This counter can be calculated at the end of DAG by simply counting the distinct
+   * ContainerIds that have been seen in TaskSchedulerManager.taskAllocated callbacks.
+   */
+  TOTAL_CONTAINERS_USED,
+
+  /*
+   * Type: # of events
+   * Number of container allocations during a DAG. This is incremented every time
+   * the containerAllocated callback is called in the TaskSchedulerContext.
+   * This counter doesn't account for initially held (launched, allocated) containers.
+   */
+  TOTAL_CONTAINER_ALLOCATION_COUNT,
+
+  /*
+   * Type: # of events
+   * Number of container launches during a DAG. This is incremented every time
+   * the containerLaunched callback is called in the ContainerLauncherContext.
+   * This counter doesn't account for initially held (launched, allocated) containers.
+   */
+  TOTAL_CONTAINER_LAUNCH_COUNT,
+
+  /*
+   * Type: # of events
+   * Number of container releases during a DAG. This is incremented every time
+   * the containerBeingReleased callback is called in the TaskSchedulerContext.
+   */
+  TOTAL_CONTAINER_RELEASE_COUNT,
+
+  /*
+   * Type: # of events
+   * Number of container reuses during a DAG. This is incremented every time
+   * the containerReused callback is called in the TaskSchedulerContext.
+   */
+  TOTAL_CONTAINER_REUSE_COUNT
 }
diff --git a/tez-api/src/main/java/org/apache/tez/serviceplugins/api/TaskScheduler.java b/tez-api/src/main/java/org/apache/tez/serviceplugins/api/TaskScheduler.java
@@ -14,6 +14,8 @@
 
 package org.apache.tez.serviceplugins.api;
 
+import java.util.List;
+
 import javax.annotation.Nullable;
 
 import org.apache.hadoop.classification.InterfaceAudience;
@@ -263,4 +265,19 @@ public abstract boolean deallocateTask(Object task, boolean taskSucceeded,
    */
   public abstract void dagComplete() throws ServicePluginException;
 
+  /**
+   * Get the number of held containers
+   * @throws ServicePluginException when the service runs into a fatal error which it cannot handle.
+   *                               This will cause the app to shutdown.
+   */
+  public abstract int getHeldContainersCount();
+
+  /**
+   * Callback to be used in the event of a container allocation.
+   */
+  protected void onContainersAllocated(List<Container> containers) {
+    for (Container container : containers) {
+      getContext().containerAllocated(container);
+    }
+  }
 }
diff --git a/tez-api/src/main/java/org/apache/tez/serviceplugins/api/TaskSchedulerContext.java b/tez-api/src/main/java/org/apache/tez/serviceplugins/api/TaskSchedulerContext.java
@@ -84,6 +84,20 @@ void taskAllocated(Object task,
                      Object appCookie,
                      Container container);
 
+  /**
+   * Indicate to the framework that a container is being allocated.
+   *
+   * @param the actual container
+   */
+  void containerAllocated(Container container);
+
+  /**
+   * Indicate to the framework that a container is being reused:
+   * there is a task assigned to an already used container.
+   *
+   * @param the actual container
+   */
+  void containerReused(Container container);
 
   /**
    * Indicate to the framework that a container has completed. This is typically used by sources
diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/ContainerLauncherContextImpl.java b/tez-dag/src/main/java/org/apache/tez/dag/app/ContainerLauncherContextImpl.java
@@ -21,6 +21,7 @@
 import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
 import org.apache.hadoop.yarn.api.records.ContainerId;
 import org.apache.tez.common.TezUtilsInternal;
+import org.apache.tez.common.counters.DAGCounter;
 import org.apache.tez.dag.api.UserPayload;
 import org.apache.tez.dag.app.dag.event.DAGAppMasterEventType;
 import org.apache.tez.dag.app.dag.event.DAGAppMasterEventUserServiceFatalError;
@@ -65,6 +66,7 @@ public ContainerLauncherContextImpl(AppContext appContext, ContainerLauncherMana
 
   @Override
   public void containerLaunched(ContainerId containerId) {
+    context.getCurrentDAG().incrementDagCounter(DAGCounter.TOTAL_CONTAINER_LAUNCH_COUNT, 1);
     context.getEventHandler().handle(
         new AMContainerEventLaunched(containerId));
     ContainerLaunchedEvent lEvt = new ContainerLaunchedEvent(
diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/DAGAppMaster.java b/tez-dag/src/main/java/org/apache/tez/dag/app/DAGAppMaster.java
@@ -121,6 +121,7 @@
 import org.apache.tez.common.TezConverterUtils;
 import org.apache.tez.common.TezUtilsInternal;
 import org.apache.tez.common.VersionInfo;
+import org.apache.tez.common.counters.DAGCounter;
 import org.apache.tez.common.counters.Limits;
 import org.apache.tez.common.security.ACLManager;
 import org.apache.tez.common.security.JobTokenIdentifier;
@@ -773,6 +774,8 @@ protected synchronized void handle(DAGAppMasterEvent event) {
       System.out.println(timeStamp + " Completed Dag: " + finishEvt.getDAGId());
       // Stop vertex services if any
       stopVertexServices(currentDAG);
+      handleUsedContainersOnDagFinish(currentDAG);
+
       if (!isSession) {
         LOG.info("Not a session, AM will unregister as DAG has completed");
         this.taskSchedulerManager.setShouldUnregisterFlag();
@@ -895,6 +898,12 @@ protected synchronized void handle(DAGAppMasterEvent event) {
     }
   }
 
+  private void handleUsedContainersOnDagFinish(DAG dag) {
+    dag.incrementDagCounter(DAGCounter.TOTAL_CONTAINERS_USED,
+        taskSchedulerManager.getContainersCountUsedByCurrentDAG());
+    taskSchedulerManager.clearContainersUsedByCurrentDAG();
+  }
+
   private void updateLoggers(DAG dag, String appender) {
     try {
       TezUtilsInternal.updateLoggers(dag.getConf(), dag.getID().toString() + appender,
@@ -2557,11 +2566,17 @@ public Void run() throws Exception {
       throw new TezUncheckedException(e);
     }
 
+    countAlreadyHeldContainers(newDAG);
     startDAGExecution(newDAG, lrDiff);
     // set state after curDag is set
     this.state = DAGAppMasterState.RUNNING;
   }
 
+  private void countAlreadyHeldContainers(DAG newDAG) {
+    newDAG.incrementDagCounter(DAGCounter.INITIAL_HELD_CONTAINERS,
+        taskSchedulerManager.getAlreadyHeldContainersCount());
+  }
+
   private void startVertexServices(DAG dag) {
     for (Vertex v : dag.getVertices().values()) {
       v.startServices();
diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/dag/DAG.java b/tez-dag/src/main/java/org/apache/tez/dag/app/dag/DAG.java
@@ -25,6 +25,7 @@
 import org.apache.hadoop.security.UserGroupInformation;
 import org.apache.hadoop.yarn.api.records.LocalResource;
 import org.apache.hadoop.yarn.event.EventHandler;
+import org.apache.tez.common.counters.DAGCounter;
 import org.apache.tez.common.counters.TezCounters;
 import org.apache.tez.dag.api.TezException;
 import org.apache.tez.dag.api.client.DAGStatusBuilder;
@@ -102,4 +103,6 @@ VertexStatusBuilder getVertexStatus(String vertexName,
    */
   @Nullable DAGScheduler getDAGScheduler();
 
+  void incrementDagCounter(DAGCounter counter, int incrValue);
+
 }
diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/DAGImpl.java b/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/DAGImpl.java
@@ -1441,6 +1441,11 @@ private void updateCpuCounters() {
     dagCounters.findCounter(DAGCounter.AM_GC_TIME_MILLIS).setValue(totalDAGGCTime);
   }
   
+  @Override
+  public void incrementDagCounter(DAGCounter counter, int incrValue) {
+    dagCounters.findCounter(counter).increment(incrValue);
+  }
+
   private DAGState finished(DAGState finalState) {
     boolean dagError = false;
     try {
diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/rm/DagAwareYarnTaskScheduler.java b/tez-dag/src/main/java/org/apache/tez/dag/app/rm/DagAwareYarnTaskScheduler.java
@@ -328,6 +328,8 @@ public void shutdown() throws Exception {
 
   @Override
   public void onContainersAllocated(List<Container> containers) {
+    super.onContainersAllocated(containers);
+
     AMState appState = getContext().getAMState();
     if (stopRequested || appState == AMState.COMPLETED) {
       LOG.info("Ignoring {} allocations since app is terminating", containers.size());
@@ -946,6 +948,9 @@ private void addTaskAssignment(TaskRequest request, HeldContainer hc) {
       assignedVertices.set(vertexIndex);
     }
     cset.add(hc);
+    if (!hc.isNew()) {
+      getContext().containerReused(hc.getContainer());
+    }
     hc.assignTask(request);
   }
 
@@ -1489,6 +1494,10 @@ Object getLastTask() {
       return lastRequest != null ? lastRequest.getTask() : null;
     }
 
+    boolean isNew() {
+      return lastRequest == null;
+    }
+
     String getMatchingLocation() {
       switch (state) {
       case MATCHING_LOCAL:
@@ -2089,4 +2098,9 @@ protected void afterExecute(Runnable r, Throwable t) {
       }
     }
   }
+
+  @Override
+  public int getHeldContainersCount() {
+    return heldContainers.size();
+  }
 }
diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/rm/LocalTaskSchedulerService.java b/tez-dag/src/main/java/org/apache/tez/dag/app/rm/LocalTaskSchedulerService.java
@@ -513,4 +513,9 @@ void preemptTask(DeallocateContainerRequest request) {
       }
     }
   }
+
+  @Override
+  public int getHeldContainersCount() {
+    return 0;
+  }
 }
diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/rm/TaskSchedulerContextImpl.java b/tez-dag/src/main/java/org/apache/tez/dag/app/rm/TaskSchedulerContextImpl.java
@@ -28,6 +28,7 @@
 import org.apache.hadoop.yarn.api.records.NodeReport;
 import org.apache.hadoop.yarn.api.records.Resource;
 import org.apache.tez.common.ContainerSignatureMatcher;
+import org.apache.tez.common.counters.DAGCounter;
 import org.apache.tez.dag.api.TezUncheckedException;
 import org.apache.tez.dag.api.UserPayload;
 import org.apache.tez.dag.app.AppContext;
@@ -69,13 +70,24 @@ public void taskAllocated(Object task, Object appCookie, Container container) {
     taskSchedulerManager.taskAllocated(schedulerId, task, appCookie, container);
   }
 
+  @Override
+  public void containerAllocated(Container container) {
+    appContext.getCurrentDAG().incrementDagCounter(DAGCounter.TOTAL_CONTAINER_ALLOCATION_COUNT, 1);
+  }
+
+  @Override
+  public void containerReused(Container container) {
+    appContext.getCurrentDAG().incrementDagCounter(DAGCounter.TOTAL_CONTAINER_REUSE_COUNT, 1);
+  }
+
   @Override
   public void containerCompleted(Object taskLastAllocated, ContainerStatus containerStatus) {
     taskSchedulerManager.containerCompleted(schedulerId, taskLastAllocated, containerStatus);
   }
 
   @Override
   public void containerBeingReleased(ContainerId containerId) {
+    appContext.getCurrentDAG().incrementDagCounter(DAGCounter.TOTAL_CONTAINER_RELEASE_COUNT, 1);
     taskSchedulerManager.containerBeingReleased(schedulerId, containerId);
   }
 
diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/rm/TaskSchedulerContextImplWrapper.java b/tez-dag/src/main/java/org/apache/tez/dag/app/rm/TaskSchedulerContextImplWrapper.java
@@ -70,6 +70,16 @@ public void taskAllocated(Object task, Object appCookie, Container container) {
         container));
   }
 
+  @Override
+  public void containerAllocated(Container container) {
+    executorService.submit(new ContainerAllocatedCallable(real, container));
+  }
+
+  @Override
+  public void containerReused(Container container) {
+    executorService.submit(new ContainerReusedCallable(real, container));
+  }
+
   @Override
   public void containerCompleted(Object taskLastAllocated,
       ContainerStatus containerStatus) {
@@ -226,6 +236,38 @@ public Void call() throws Exception {
     }
   }
 
+  static class ContainerAllocatedCallable extends TaskSchedulerContextCallbackBase
+      implements Callable<Void> {
+    private final Container container;
+
+    public ContainerAllocatedCallable(TaskSchedulerContext app, Container container) {
+      super(app);
+      this.container = container;
+    }
+
+    @Override
+    public Void call() throws Exception {
+      app.containerAllocated(container);
+      return null;
+    }
+  }
+
+  static class ContainerReusedCallable extends TaskSchedulerContextCallbackBase
+      implements Callable<Void> {
+    private final Container container;
+
+    public ContainerReusedCallable(TaskSchedulerContext app, Container container) {
+      super(app);
+      this.container = container;
+    }
+
+    @Override
+    public Void call() throws Exception {
+      app.containerReused(container);
+      return null;
+    }
+  }
+
   static class ContainerCompletedCallable extends TaskSchedulerContextCallbackBase
       implements Callable<Void> {
 
diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/rm/TaskSchedulerManager.java b/tez-dag/src/main/java/org/apache/tez/dag/app/rm/TaskSchedulerManager.java
@@ -20,6 +20,7 @@
 
 import java.net.InetSocketAddress;
 import java.nio.ByteBuffer;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.BlockingQueue;
@@ -29,6 +30,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.Objects;
+import java.util.Set;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
@@ -122,6 +124,7 @@ public class TaskSchedulerManager extends AbstractService implements
   private final NamedEntityDescriptor[] taskSchedulerDescriptors;
   protected final TaskSchedulerWrapper[] taskSchedulers;
   protected final ServicePluginLifecycleAbstractService []taskSchedulerServiceWrappers;
+  protected final Set<ContainerId> containersUsedByCurrentDAG = new HashSet<>();
 
   // Single executor service shared by all Schedulers for context callbacks
   @VisibleForTesting
@@ -746,7 +749,7 @@ public synchronized void taskAllocated(int schedulerId, Object task,
       sendEvent(new AMNodeEventContainerAllocated(container
           .getNodeId(), schedulerId, container.getId()));
     }
-
+    containersUsedByCurrentDAG.add(containerId);
 
     TaskAttempt taskAttempt = event.getTaskAttempt();
     // TODO - perhaps check if the task still needs this container
@@ -951,6 +954,14 @@ public void dagCompleted() {
     }
   }
 
+  public int getAlreadyHeldContainersCount() {
+    int count = 0;
+    for (int i = 0 ; i < taskSchedulers.length ; i++) {
+      count += taskSchedulers[i].getTaskScheduler().getHeldContainersCount();
+    }
+    return count;
+  }
+
   public void dagSubmitted() {
     // Nothing to do right now. Indicates that a new DAG has been submitted and
     // the context has updated information.
@@ -1084,4 +1095,11 @@ public TaskScheduler getTaskScheduler(int taskSchedulerIndex) {
     return taskSchedulers[taskSchedulerIndex].getTaskScheduler();
   }
 
+  public int getContainersCountUsedByCurrentDAG() {
+    return containersUsedByCurrentDAG.size();
+  }
+
+  public void clearContainersUsedByCurrentDAG() {
+    containersUsedByCurrentDAG.clear();
+  }
 }
diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/rm/YarnTaskSchedulerService.java b/tez-dag/src/main/java/org/apache/tez/dag/app/rm/YarnTaskSchedulerService.java
diff --git a/tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestTaskSchedulerHelpers.java b/tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestTaskSchedulerHelpers.java
diff --git a/tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestTaskSchedulerManager.java b/tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestTaskSchedulerManager.java
diff --git a/tez-ext-service-tests/src/test/java/org/apache/tez/dag/app/rm/TezTestServiceTaskSchedulerService.java b/tez-ext-service-tests/src/test/java/org/apache/tez/dag/app/rm/TezTestServiceTaskSchedulerService.java
diff --git a/tez-ext-service-tests/src/test/java/org/apache/tez/dag/app/rm/TezTestServiceTaskSchedulerServiceWithErrors.java b/tez-ext-service-tests/src/test/java/org/apache/tez/dag/app/rm/TezTestServiceTaskSchedulerServiceWithErrors.java

Original file line number	Diff line number	Diff line change
`@@ -328,6 +328,8 @@ public void shutdown() throws Exception {`
`328`	`328`
`329`	`329`	`@Override`
`330`	`330`	`public void onContainersAllocated(List<Container> containers) {`
	`331`	`+ super.onContainersAllocated(containers);`
	`332`	`+`
`331`	`333`	`AMState appState = getContext().getAMState();`
`332`	`334`	`if (stopRequested \|\| appState == AMState.COMPLETED) {`
`333`	`335`	`LOG.info("Ignoring {} allocations since app is terminating", containers.size());`
`@@ -946,6 +948,9 @@ private void addTaskAssignment(TaskRequest request, HeldContainer hc) {`
`946`	`948`	`assignedVertices.set(vertexIndex);`
`947`	`949`	`}`
`948`	`950`	`cset.add(hc);`
	`951`	`+ if (!hc.isNew()) {`
	`952`	`+ getContext().containerReused(hc.getContainer());`
	`953`	`+ }`
`949`	`954`	`hc.assignTask(request);`
`950`	`955`	`}`
`951`	`956`
`@@ -1489,6 +1494,10 @@ Object getLastTask() {`
`1489`	`1494`	`return lastRequest != null ? lastRequest.getTask() : null;`
`1490`	`1495`	`}`
`1491`	`1496`
	`1497`	`+ boolean isNew() {`
	`1498`	`+ return lastRequest == null;`
	`1499`	`+ }`
	`1500`	`+`
`1492`	`1501`	`String getMatchingLocation() {`
`1493`	`1502`	`switch (state) {`
`1494`	`1503`	`case MATCHING_LOCAL:`
`@@ -2089,4 +2098,9 @@ protected void afterExecute(Runnable r, Throwable t) {`
`2089`	`2098`	`}`
`2090`	`2099`	`}`
`2091`	`2100`	`}`
	`2101`	`+`
	`2102`	`+ @Override`
	`2103`	`+ public int getHeldContainersCount() {`
	`2104`	`+ return heldContainers.size();`
	`2105`	`+ }`
`2092`	`2106`	`}`
Original file line number	Diff line number	Diff line change
`@@ -513,4 +513,9 @@ void preemptTask(DeallocateContainerRequest request) {`
`513`	`513`	`}`
`514`	`514`	`}`
`515`	`515`	`}`
	`516`	`+`
	`517`	`+ @Override`
	`518`	`+ public int getHeldContainersCount() {`
	`519`	`+ return 0;`
	`520`	`+ }`
`516`	`521`	`}`