[SPARK-26329][CORE] Faster polling of executor memory metrics.

## What changes were proposed in this pull request?

Prior to this change, in an executor, on each heartbeat, memory metrics are polled and sent in the heartbeat. The heartbeat interval is 10s by default. With this change, in an executor, memory metrics can optionally be polled in a separate poller at a shorter interval.

For each executor, we use a map of (stageId, stageAttemptId) to (count of running tasks, executor metric peaks) to track what stages are active as well as the per-stage memory metric peaks. When polling the executor memory metrics, we attribute the memory to the active stage(s), and update the peaks. In a heartbeat, we send the per-stage peaks (for stages active at that time), and then reset the peaks. The semantics would be that the per-stage peaks sent in each heartbeat are the peaks since the last heartbeat.

We also keep a map of taskId to memory metric peaks. This tracks the metric peaks during the lifetime of the task. The polling thread updates this as well. At end of a task, we send the peak metric values in the task result. In case of task failure, we send the peak metric values in the `TaskFailedReason`.

We continue to do the stage-level aggregation in the EventLoggingListener.

For the driver, we still only poll on heartbeats. What the driver sends will be the current values of the metrics in the driver at the time of the heartbeat. This is semantically the same as before.

## How was this patch tested?

Unit tests. Manually tested applications on an actual system and checked the event logs; the metrics appear in the SparkListenerTaskEnd and SparkListenerStageExecutorMetrics events.

Closes #23767 from wypoon/wypoon_SPARK-26329.

Authored-by: Wing Yew Poon <[email protected]>
Signed-off-by: Imran Rashid <[email protected]>

Author: wypoon

core/src/main/scala/org/apache/spark/HeartbeatReceiver.scala

@@ -19,7 +19,7 @@ package org.apache.spark
 
 import java.util.concurrent.{ScheduledFuture, TimeUnit}
 
-import scala.collection.mutable
+import scala.collection.mutable.{HashMap, Map}
 import scala.concurrent.Future
 
 import org.apache.spark.executor.ExecutorMetrics
@@ -38,9 +38,11 @@ import org.apache.spark.util._
  */
 private[spark] case class Heartbeat(
     executorId: String,
-    accumUpdates: Array[(Long, Seq[AccumulatorV2[_, _]])], // taskId -> accumulator updates
+    // taskId -> accumulator updates
+    accumUpdates: Array[(Long, Seq[AccumulatorV2[_, _]])],
     blockManagerId: BlockManagerId,
-    executorUpdates: ExecutorMetrics) // executor level updates
+    // (stageId, stageAttemptId) -> executor metric peaks
+    executorUpdates: Map[(Int, Int), ExecutorMetrics])
 
 /**
  * An event that SparkContext uses to notify HeartbeatReceiver that SparkContext.taskScheduler is
@@ -73,7 +75,7 @@ private[spark] class HeartbeatReceiver(sc: SparkContext, clock: Clock)
   private[spark] var scheduler: TaskScheduler = null
 
   // executor ID -> timestamp of when the last heartbeat from this executor was received
-  private val executorLastSeen = new mutable.HashMap[String, Long]
+  private val executorLastSeen = new HashMap[String, Long]
 
   private val executorTimeoutMs = sc.conf.get(config.STORAGE_BLOCKMANAGER_SLAVE_TIMEOUT)
 
@@ -120,14 +122,14 @@ private[spark] class HeartbeatReceiver(sc: SparkContext, clock: Clock)
       context.reply(true)
 
     // Messages received from executors
-    case heartbeat @ Heartbeat(executorId, accumUpdates, blockManagerId, executorMetrics) =>
+    case heartbeat @ Heartbeat(executorId, accumUpdates, blockManagerId, executorUpdates) =>
       if (scheduler != null) {
         if (executorLastSeen.contains(executorId)) {
           executorLastSeen(executorId) = clock.getTimeMillis()
           eventLoopThread.submit(new Runnable {
             override def run(): Unit = Utils.tryLogNonFatalError {
               val unknownExecutor = !scheduler.executorHeartbeatReceived(
-                executorId, accumUpdates, blockManagerId, executorMetrics)
+                executorId, accumUpdates, blockManagerId, executorUpdates)
               val response = HeartbeatResponse(reregisterBlockManager = unknownExecutor)
               context.reply(response)
             }

core/src/main/scala/org/apache/spark/Heartbeater.scala

@@ -19,23 +19,18 @@ package org.apache.spark
 
 import java.util.concurrent.TimeUnit
 
-import org.apache.spark.executor.ExecutorMetrics
 import org.apache.spark.internal.Logging
-import org.apache.spark.memory.MemoryManager
-import org.apache.spark.metrics.ExecutorMetricType
 import org.apache.spark.util.{ThreadUtils, Utils}
 
 /**
  * Creates a heartbeat thread which will call the specified reportHeartbeat function at
  * intervals of intervalMs.
  *
- * @param memoryManager the memory manager for execution and storage memory.
  * @param reportHeartbeat the heartbeat reporting function to call.
  * @param name the thread name for the heartbeater.
  * @param intervalMs the interval between heartbeats.
  */
 private[spark] class Heartbeater(
-    memoryManager: MemoryManager,
     reportHeartbeat: () => Unit,
     name: String,
     intervalMs: Long) extends Logging {
@@ -58,21 +53,4 @@ private[spark] class Heartbeater(
     heartbeater.shutdown()
     heartbeater.awaitTermination(10, TimeUnit.SECONDS)
   }
-
-  /**
-   * Get the current executor level metrics. These are returned as an array, with the index
-   * determined by ExecutorMetricType.metricToOffset
-   */
-  def getCurrentMetrics(): ExecutorMetrics = {
-
-    val metrics = new Array[Long](ExecutorMetricType.numMetrics)
-    var offset = 0
-    ExecutorMetricType.metricGetters.foreach { metric =>
-      val newMetrics = metric.getMetricValues(memoryManager)
-      Array.copy(newMetrics, 0, metrics, offset, newMetrics.size)
-      offset += newMetrics.length
-    }
-    new ExecutorMetrics(metrics)
-  }
 }
-

core/src/main/scala/org/apache/spark/SparkContext.scala

@@ -42,6 +42,7 @@ import org.apache.hadoop.mapreduce.lib.input.{FileInputFormat => NewFileInputFor
 import org.apache.spark.annotation.DeveloperApi
 import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.deploy.{LocalSparkCluster, SparkHadoopUtil}
+import org.apache.spark.executor.ExecutorMetrics
 import org.apache.spark.input.{FixedLengthBinaryInputFormat, PortableDataStream, StreamInputFormat, WholeTextFileInputFormat}
 import org.apache.spark.internal.Logging
 import org.apache.spark.internal.config._
@@ -516,7 +517,7 @@ class SparkContext(config: SparkConf) extends Logging {
     _heartbeatReceiver.ask[Boolean](TaskSchedulerIsSet)
 
     // create and start the heartbeater for collecting memory metrics
-    _heartbeater = new Heartbeater(env.memoryManager,
+    _heartbeater = new Heartbeater(
       () => SparkContext.this.reportHeartBeat(),
       "driver-heartbeater",
       conf.get(EXECUTOR_HEARTBEAT_INTERVAL))
@@ -2425,10 +2426,13 @@ class SparkContext(config: SparkConf) extends Logging {
 
   /** Reports heartbeat metrics for the driver. */
   private def reportHeartBeat(): Unit = {
-    val driverUpdates = _heartbeater.getCurrentMetrics()
+    val currentMetrics = ExecutorMetrics.getCurrentMetrics(env.memoryManager)
+    val driverUpdates = new HashMap[(Int, Int), ExecutorMetrics]
+    // In the driver, we do not track per-stage metrics, so use a dummy stage for the key
+    driverUpdates.put(EventLoggingListener.DRIVER_STAGE_KEY, new ExecutorMetrics(currentMetrics))
     val accumUpdates = new Array[(Long, Int, Int, Seq[AccumulableInfo])](0)
     listenerBus.post(SparkListenerExecutorMetricsUpdate("driver", accumUpdates,
-      Some(driverUpdates)))
+      driverUpdates))
   }
 
   // In order to prevent multiple SparkContexts from being active at the same time, mark this

core/src/main/scala/org/apache/spark/TaskEndReason.scala

@@ -128,7 +128,8 @@ case class ExceptionFailure(
     fullStackTrace: String,
     private val exceptionWrapper: Option[ThrowableSerializationWrapper],
     accumUpdates: Seq[AccumulableInfo] = Seq.empty,
-    private[spark] var accums: Seq[AccumulatorV2[_, _]] = Nil)
+    private[spark] var accums: Seq[AccumulatorV2[_, _]] = Nil,
+    private[spark] var metricPeaks: Seq[Long] = Seq.empty)
   extends TaskFailedReason {
 
   /**
@@ -153,6 +154,11 @@ case class ExceptionFailure(
     this
   }
 
+  private[spark] def withMetricPeaks(metricPeaks: Seq[Long]): ExceptionFailure = {
+    this.metricPeaks = metricPeaks
+    this
+  }
+
   def exception: Option[Throwable] = exceptionWrapper.flatMap(w => Option(w.exception))
 
   override def toErrorString: String =
@@ -215,7 +221,8 @@ case object TaskResultLost extends TaskFailedReason {
 case class TaskKilled(
     reason: String,
     accumUpdates: Seq[AccumulableInfo] = Seq.empty,
-    private[spark] val accums: Seq[AccumulatorV2[_, _]] = Nil)
+    private[spark] val accums: Seq[AccumulatorV2[_, _]] = Nil,
+    metricPeaks: Seq[Long] = Seq.empty)
   extends TaskFailedReason {
 
   override def toErrorString: String = s"TaskKilled ($reason)"

core/src/main/scala/org/apache/spark/executor/Executor.scala

@@ -27,7 +27,7 @@ import java.util.concurrent._
 import javax.annotation.concurrent.GuardedBy
 
 import scala.collection.JavaConverters._
-import scala.collection.mutable.{ArrayBuffer, HashMap, Map}
+import scala.collection.mutable.{ArrayBuffer, HashMap, Map, WrappedArray}
 import scala.concurrent.duration._
 import scala.util.control.NonFatal
 
@@ -188,9 +188,20 @@ private[spark] class Executor(
    */
   private val HEARTBEAT_INTERVAL_MS = conf.get(EXECUTOR_HEARTBEAT_INTERVAL)
 
+  /**
+   * Interval to poll for executor metrics, in milliseconds
+   */
+  private val METRICS_POLLING_INTERVAL_MS = conf.get(EXECUTOR_METRICS_POLLING_INTERVAL)
+
+  private val pollOnHeartbeat = if (METRICS_POLLING_INTERVAL_MS > 0) false else true
+
+  // Poller for the memory metrics. Visible for testing.
+  private[executor] val metricsPoller = new ExecutorMetricsPoller(
+    env.memoryManager,
+    METRICS_POLLING_INTERVAL_MS)
+
   // Executor for the heartbeat task.
   private val heartbeater = new Heartbeater(
-    env.memoryManager,
     () => Executor.this.reportHeartBeat(),
     "executor-heartbeater",
     HEARTBEAT_INTERVAL_MS)
@@ -207,6 +218,8 @@ private[spark] class Executor(
 
   heartbeater.start()
 
+  metricsPoller.start()
+
   private[executor] def numRunningTasks: Int = runningTasks.size()
 
   def launchTask(context: ExecutorBackend, taskDescription: TaskDescription): Unit = {
@@ -254,12 +267,18 @@ private[spark] class Executor(
 
   def stop(): Unit = {
     env.metricsSystem.report()
+    try {
+      metricsPoller.stop()
+    } catch {
+      case NonFatal(e) =>
+        logWarning("Unable to stop executor metrics poller", e)
+    }
     try {
       heartbeater.stop()
     } catch {
       case NonFatal(e) =>
         logWarning("Unable to stop heartbeater", e)
-     }
+    }
     threadPool.shutdown()
 
     // Notify plugins that executor is shutting down so they can terminate cleanly
@@ -380,6 +399,7 @@ private[spark] class Executor(
       var taskStartTimeNs: Long = 0
       var taskStartCpu: Long = 0
       startGCTime = computeTotalGcTime()
+      var taskStarted: Boolean = false
 
       try {
         // Must be set before updateDependencies() is called, in case fetching dependencies
@@ -412,6 +432,9 @@ private[spark] class Executor(
           env.mapOutputTracker.asInstanceOf[MapOutputTrackerWorker].updateEpoch(task.epoch)
         }
 
+        metricsPoller.onTaskStart(taskId, task.stageId, task.stageAttemptId)
+        taskStarted = true
+
         // Run the actual task and measure its runtime.
         taskStartTimeNs = System.nanoTime()
         taskStartCpu = if (threadMXBean.isCurrentThreadCpuTimeSupported) {
@@ -529,8 +552,9 @@ private[spark] class Executor(
 
         // Note: accumulator updates must be collected after TaskMetrics is updated
         val accumUpdates = task.collectAccumulatorUpdates()
+        val metricPeaks = metricsPoller.getTaskMetricPeaks(taskId)
         // TODO: do not serialize value twice
-        val directResult = new DirectTaskResult(valueBytes, accumUpdates)
+        val directResult = new DirectTaskResult(valueBytes, accumUpdates, metricPeaks)
         val serializedDirectResult = ser.serialize(directResult)
         val resultSize = serializedDirectResult.limit()
 
@@ -559,13 +583,15 @@ private[spark] class Executor(
         executorSource.SUCCEEDED_TASKS.inc(1L)
         setTaskFinishedAndClearInterruptStatus()
         execBackend.statusUpdate(taskId, TaskState.FINISHED, serializedResult)
-
       } catch {
         case t: TaskKilledException =>
           logInfo(s"Executor killed $taskName (TID $taskId), reason: ${t.reason}")
 
           val (accums, accUpdates) = collectAccumulatorsAndResetStatusOnFailure(taskStartTimeNs)
-          val serializedTK = ser.serialize(TaskKilled(t.reason, accUpdates, accums))
+          // Here and below, put task metric peaks in a WrappedArray to expose them as a Seq
+          // without requiring a copy.
+          val metricPeaks = WrappedArray.make(metricsPoller.getTaskMetricPeaks(taskId))
+          val serializedTK = ser.serialize(TaskKilled(t.reason, accUpdates, accums, metricPeaks))
           execBackend.statusUpdate(taskId, TaskState.KILLED, serializedTK)
 
         case _: InterruptedException | NonFatal(_) if
@@ -574,7 +600,8 @@ private[spark] class Executor(
           logInfo(s"Executor interrupted and killed $taskName (TID $taskId), reason: $killReason")
 
           val (accums, accUpdates) = collectAccumulatorsAndResetStatusOnFailure(taskStartTimeNs)
-          val serializedTK = ser.serialize(TaskKilled(killReason, accUpdates, accums))
+          val metricPeaks = WrappedArray.make(metricsPoller.getTaskMetricPeaks(taskId))
+          val serializedTK = ser.serialize(TaskKilled(killReason, accUpdates, accums, metricPeaks))
           execBackend.statusUpdate(taskId, TaskState.KILLED, serializedTK)
 
         case t: Throwable if hasFetchFailure && !Utils.isFatalError(t) =>
@@ -609,14 +636,19 @@ private[spark] class Executor(
           // instead of an app issue).
           if (!ShutdownHookManager.inShutdown()) {
             val (accums, accUpdates) = collectAccumulatorsAndResetStatusOnFailure(taskStartTimeNs)
+            val metricPeaks = WrappedArray.make(metricsPoller.getTaskMetricPeaks(taskId))
 
             val serializedTaskEndReason = {
               try {
-                ser.serialize(new ExceptionFailure(t, accUpdates).withAccums(accums))
+                val ef = new ExceptionFailure(t, accUpdates).withAccums(accums)
+                  .withMetricPeaks(metricPeaks)
+                ser.serialize(ef)
               } catch {
                 case _: NotSerializableException =>
                   // t is not serializable so just send the stacktrace
-                  ser.serialize(new ExceptionFailure(t, accUpdates, false).withAccums(accums))
+                  val ef = new ExceptionFailure(t, accUpdates, false).withAccums(accums)
+                    .withMetricPeaks(metricPeaks)
+                  ser.serialize(ef)
               }
             }
             setTaskFinishedAndClearInterruptStatus()
@@ -632,6 +664,11 @@ private[spark] class Executor(
           }
       } finally {
         runningTasks.remove(taskId)
+        if (taskStarted) {
+          // This means the task was successfully deserialized, its stageId and stageAttemptId
+          // are known, and metricsPoller.onTaskStart was called.
+          metricsPoller.onTaskCompletion(taskId, task.stageId, task.stageAttemptId)
+        }
       }
     }
 
@@ -848,8 +885,11 @@ private[spark] class Executor(
     val accumUpdates = new ArrayBuffer[(Long, Seq[AccumulatorV2[_, _]])]()
     val curGCTime = computeTotalGcTime()
 
-    // get executor level memory metrics
-    val executorUpdates = heartbeater.getCurrentMetrics()
+    if (pollOnHeartbeat) {
+      metricsPoller.poll()
+    }
+
+    val executorUpdates = metricsPoller.getExecutorUpdates()
 
     for (taskRunner <- runningTasks.values().asScala) {
       if (taskRunner.task != null) {

core/src/main/scala/org/apache/spark/executor/ExecutorMetrics.scala

@@ -16,7 +16,10 @@
  */
 package org.apache.spark.executor
 
+import java.util.concurrent.atomic.AtomicLongArray
+
 import org.apache.spark.annotation.DeveloperApi
+import org.apache.spark.memory.MemoryManager
 import org.apache.spark.metrics.ExecutorMetricType
 
 /**
@@ -46,14 +49,21 @@ class ExecutorMetrics private[spark] extends Serializable {
     Array.copy(metrics, 0, this.metrics, 0, Math.min(metrics.size, this.metrics.size))
   }
 
+  private[spark] def this(metrics: AtomicLongArray) {
+    this()
+    ExecutorMetricType.metricToOffset.foreach { case (_, i) =>
+      this.metrics(i) = metrics.get(i)
+    }
+  }
+
   /**
-   * Constructor: create the ExecutorMetrics with using a given map.
+   * Constructor: create the ExecutorMetrics using a given map.
    *
    * @param executorMetrics map of executor metric name to value
    */
   private[spark] def this(executorMetrics: Map[String, Long]) {
     this()
-    ExecutorMetricType.metricToOffset.foreach { case(name, idx) =>
+    ExecutorMetricType.metricToOffset.foreach { case (name, idx) =>
       metrics(idx) = executorMetrics.getOrElse(name, 0L)
     }
   }
@@ -76,3 +86,24 @@ class ExecutorMetrics private[spark] extends Serializable {
     updated
   }
 }
+
+private[spark] object ExecutorMetrics {
+
+  /**
+   * Get the current executor metrics. These are returned as an array, with the index
+   * determined by ExecutorMetricType.metricToOffset.
+   *
+   * @param memoryManager the memory manager for execution and storage memory
+   * @return the values of the metrics
+   */
+  def getCurrentMetrics(memoryManager: MemoryManager): Array[Long] = {
+    val currentMetrics = new Array[Long](ExecutorMetricType.numMetrics)
+    var offset = 0
+    ExecutorMetricType.metricGetters.foreach { metricType =>
+      val metricValues = metricType.getMetricValues(memoryManager)
+      Array.copy(metricValues, 0, currentMetrics, offset, metricValues.length)
+      offset += metricValues.length
+    }
+    currentMetrics
+  }
+}