TEZ-4601: New counters in task scheduler: request times sum, max, average

Author: abstractdog

tez-api/src/main/java/org/apache/tez/common/counters/DAGCounter.java

@@ -122,5 +122,18 @@ public enum DAGCounter {
    * task assignments). This is typically exposed by a resource manager
    * client.
    */
-  NODE_TOTAL_COUNT
+  NODE_TOTAL_COUNT,
+
+  /*
+   * The maximum amount of time a task spent waiting as a task request before being scheduled.
+   */
+  TASK_SCHEDULER_MAX_PENDING_TIME_MS,
+  /*
+   * The total accumulated time that all tasks spent waiting as task requests.
+   */
+  TASK_SCHEDULER_SUM_PENDING_TIME_MS,
+  /*
+   * The average time tasks spent waiting as task requests before being scheduled.
+   */
+  TASK_SCHEDULER_AVG_PENDING_TIME_MS
 }

tez-api/src/main/java/org/apache/tez/serviceplugins/api/TaskScheduler.java

@@ -26,6 +26,10 @@
 import org.apache.hadoop.yarn.api.records.Priority;
 import org.apache.hadoop.yarn.api.records.Resource;
 import org.apache.tez.common.ServicePluginLifecycle;
+import org.apache.tez.common.counters.AbstractCounters;
+import org.apache.tez.common.counters.CounterGroup;
+import org.apache.tez.common.counters.TezCounter;
+import org.apache.tez.common.counters.TezCounters;
 
 /**
  * This class represents the API for a custom TaskScheduler which can be run within the Tez AM.
@@ -54,7 +58,6 @@ public enum SchedulerTaskState {
   public TaskScheduler(TaskSchedulerContext taskSchedulerContext) {
     this.taskSchedulerContext = taskSchedulerContext;
   }
-
   /**
    * An entry point for initialization.
    * Order of service setup. Constructor, initialize(), start() - when starting a service.
@@ -280,4 +283,12 @@ protected void onContainersAllocated(List<Container> containers) {
       getContext().containerAllocated(container);
     }
   }
+
+  /**
+   * Collects DAG-level counters from the TaskScheduler, which is then aggregated by the DAG implementation.
+   * @return null by default, handled in upper layers
+   */
+  public TaskSchedulerStatistics getStatistics() {
+    return null;
+  }
 }

tez-api/src/main/java/org/apache/tez/serviceplugins/api/TaskSchedulerStatistics.java

@@ -0,0 +1,112 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.tez.serviceplugins.api;
+
+import org.apache.hadoop.util.Time;
+import org.apache.tez.common.counters.DAGCounter;
+import org.apache.tez.common.counters.TezCounters;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Statistics aggregator for task scheduler requests.
+ */
+public class TaskSchedulerStatistics {
+  private static final Logger LOG = LoggerFactory.getLogger(TaskSchedulerStatistics.class);
+
+  /*
+   * Tracking task allocation pending times.
+   */
+  private int maxPendingTime = 0;
+  private int sumPendingTime = 0;
+  private int averagePendingTime = 0;
+  private int pendingTaskRequestSamples = 0;
+
+  /**
+   * Called by task schedulers when a request is removed.
+   */
+  public void trackRequestPendingTime(TaskRequestData request) {
+    if (request == null) {
+      LOG.info("Got null TaskRequestData, ignoring (it's fine in case of early shutdowns)");
+      return;
+    }
+    long requestPendingTime = Time.now() - request.getCreatedTime();
+    LOG.debug("Tracking request pending time: {}", requestPendingTime);
+
+    sumPendingTime += (int) requestPendingTime;
+    pendingTaskRequestSamples += 1;
+    maxPendingTime = Math.max(maxPendingTime, (int) requestPendingTime);
+  }
+
+  /**
+   * Calculates some derived statistics.
+   * @return this instance
+   */
+  protected TaskSchedulerStatistics aggregate() {
+    // it's fine to lose float precision here, we're interested in average of milliseconds
+    this.averagePendingTime = pendingTaskRequestSamples > 0 ? sumPendingTime / pendingTaskRequestSamples : 0;
+    return this;
+  }
+
+  /**
+   * Adds the stats from another aggregator to this one and calculates the derived fields by calling aggregate().
+   * @return this instance with ready-to-use data
+   */
+  public TaskSchedulerStatistics add(TaskSchedulerStatistics other) {
+    if (other == null || other.pendingTaskRequestSamples == 0) {
+      return aggregate();
+    }
+    this.maxPendingTime = Math.max(maxPendingTime, other.maxPendingTime);
+    this.sumPendingTime += other.sumPendingTime;
+    this.pendingTaskRequestSamples += other.pendingTaskRequestSamples;
+    return aggregate();
+  }
+
+  public TezCounters getCounters() {
+    TezCounters counters = new TezCounters();
+    LOG.info("Getting counters from statistics: {}", this.statsToString());
+    // prevent filling the counters with useless 0 values if any (e.g. in case of unused TaskScheduler)
+    if (pendingTaskRequestSamples != 0) {
+      counters.findCounter(DAGCounter.TASK_SCHEDULER_MAX_PENDING_TIME_MS).setValue(maxPendingTime);
+      counters.findCounter(DAGCounter.TASK_SCHEDULER_SUM_PENDING_TIME_MS).setValue(sumPendingTime);
+      counters.findCounter(DAGCounter.TASK_SCHEDULER_AVG_PENDING_TIME_MS).setValue(averagePendingTime);
+    }
+    return counters;
+  }
+
+  public TaskSchedulerStatistics clear() {
+    averagePendingTime = 0;
+    sumPendingTime = 0;
+    pendingTaskRequestSamples = 0;
+    maxPendingTime = 0;
+    return this;
+  }
+
+  /**
+   * Classes implementing this interface tells basic characteristics about task requests they encapsulate.
+   */
+  public interface TaskRequestData {
+    // The time when the request was created
+    long getCreatedTime();
+  }
+
+  public String statsToString(){
+    return String.format("[pending: {max: %d, sum: %d, average: %d, samples: %d}]", maxPendingTime, sumPendingTime,
+        averagePendingTime, pendingTaskRequestSamples);
+  }
+}

tez-dag/src/main/java/org/apache/tez/dag/app/DAGAppMaster.java

@@ -73,6 +73,7 @@
 import org.apache.tez.client.TezClientUtils;
 import org.apache.tez.common.ReflectionUtils;
 import org.apache.tez.common.TezUtils;
+import org.apache.tez.common.counters.TezCounters;
 import org.apache.tez.dag.api.NamedEntityDescriptor;
 import org.apache.tez.dag.api.SessionNotRunning;
 import org.apache.tez.dag.api.UserPayload;

tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/DAGImpl.java

@@ -2591,5 +2591,7 @@ private void updateCounters() {
     setDagCounter(DAGCounter.TOTAL_CONTAINERS_USED, containersUsedByCurrentDAG.size());
     setDagCounter(DAGCounter.NODE_USED_COUNT, nodesUsedByCurrentDAG.size());
     setDagCounter(DAGCounter.NODE_TOTAL_COUNT, appContext.getTaskScheduler().getNumClusterNodes(true));
+
+    dagCounters.incrAllCounters(appContext.getTaskScheduler().getCounters());
   }
 }

tez-dag/src/main/java/org/apache/tez/dag/app/rm/DagAwareYarnTaskScheduler.java

@@ -43,13 +43,17 @@
 import org.apache.hadoop.yarn.util.resource.Resources;
 import org.apache.tez.common.ContainerSignatureMatcher;
 import org.apache.tez.common.TezUtils;
+import org.apache.tez.common.counters.AbstractCounters;
+import org.apache.tez.common.counters.CounterGroup;
+import org.apache.tez.common.counters.TezCounter;
 import org.apache.tez.dag.api.TezConfiguration;
 import org.apache.tez.dag.app.dag.TaskAttempt;
 import org.apache.tez.serviceplugins.api.DagInfo;
 import org.apache.tez.serviceplugins.api.TaskAttemptEndReason;
 import org.apache.tez.serviceplugins.api.TaskScheduler;
 import org.apache.tez.serviceplugins.api.TaskSchedulerContext;
 import org.apache.tez.serviceplugins.api.TaskSchedulerContext.AMState;
+import org.apache.tez.serviceplugins.api.TaskSchedulerStatistics;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -157,6 +161,8 @@ public class DagAwareYarnTaskScheduler extends TaskScheduler
   private int lastPreemptionHeartbeat = 0;
   private long preemptionMaxWaitTime;
 
+  private final TaskSchedulerStatistics taskSchedulerStatistics = new TaskSchedulerStatistics();
+
   public DagAwareYarnTaskScheduler(TaskSchedulerContext taskSchedulerContext) {
     super(taskSchedulerContext);
     signatureMatcher = taskSchedulerContext.getContainerSignatureMatcher();
@@ -595,6 +601,7 @@ private void informAppAboutAssignment(TaskRequest request, Container container)
           request.getCookie());
     } else {
       getContext().taskAllocated(request.getTask(), request.getCookie(), container);
+      taskSchedulerStatistics.trackRequestPendingTime(request);
     }
   }
 
@@ -1162,6 +1169,7 @@ public synchronized void dagComplete() {
       hc.resetMatchingLevel();
     }
     vertexDescendants = null;
+    taskSchedulerStatistics.clear();
   }
 
   @GuardedBy("this")
@@ -1265,12 +1273,14 @@ public void updateBlacklist(List<String> additions, List<String> removals) {
   /**
    * A utility class to track a task allocation.
    */
-  static class TaskRequest extends AMRMClient.ContainerRequest {
+  static class TaskRequest extends AMRMClient.ContainerRequest
+      implements TaskSchedulerStatistics.TaskRequestData {
     final Object task;
     final int vertexIndex;
     final Object signature;
     final Object cookie;
     final ContainerId affinityContainerId;
+    final long created;
 
     TaskRequest(Object task, int vertexIndex, Resource capability, String[] hosts, String[] racks,
         Priority priority, Object signature, Object cookie) {
@@ -1285,6 +1295,7 @@ static class TaskRequest extends AMRMClient.ContainerRequest {
       this.signature = signature;
       this.cookie = cookie;
       this.affinityContainerId = affinityContainerId;
+      this.created = Time.now();
     }
 
     Object getTask() {
@@ -1313,6 +1324,11 @@ boolean hasLocality() {
       List<String> racks = getRacks();
       return (nodes != null && !nodes.isEmpty()) || (racks != null && !racks.isEmpty());
     }
+
+    @Override
+    public long getCreatedTime() {
+      return created;
+    }
   }
 
   private enum HeldContainerState {
@@ -2103,4 +2119,9 @@ protected void afterExecute(Runnable r, Throwable t) {
   public int getHeldContainersCount() {
     return heldContainers.size();
   }
+
+  @Override
+  public TaskSchedulerStatistics getStatistics() {
+    return taskSchedulerStatistics;
+  }
 }