[SPARK-23431][CORE] Expose the new executor memory metrics at the stage level

core/src/main/scala/org/apache/spark/status/AppStatusListener.scala

@@ -790,6 +790,21 @@ private[spark] class AppStatusListener(
       }
     }
 
+    // update stage executor metrics
+    event.executorUpdates.foreach { updates =>
+      val activeStages = liveStages.values().asScala
+        .filter(_.status == v1.StageStatus.ACTIVE)
+      activeStages.foreach { stage =>
+        if (stage.peakExecutorMetrics.compareAndUpdatePeakValues(updates)) {
+          maybeUpdate(stage, now)
+        }
+        val executorSummary = stage.executorSummary(event.execId)
+        if (executorSummary.peakExecutorMetrics.compareAndUpdatePeakValues(updates)) {
+          maybeUpdate(executorSummary, now)
+        }
+      }
+    }
+
     // check if there is a new peak value for any of the executor level memory metrics
     // for the live UI. SparkListenerExecutorMetricsUpdate events are only processed
     // for the live UI.
@@ -802,15 +817,24 @@ private[spark] class AppStatusListener(
     }
   }
 
-  override def onStageExecutorMetrics(executorMetrics: SparkListenerStageExecutorMetrics): Unit = {
+  override def onStageExecutorMetrics(event: SparkListenerStageExecutorMetrics): Unit = {
     val now = System.nanoTime()
-
+    Option(liveStages.get((event.stageId, event.stageAttemptId)))
+      .foreach { stage =>
+      if (stage.peakExecutorMetrics.compareAndUpdatePeakValues(event.executorMetrics)) {
+        update(stage, now)
+      }
+      val esummary = stage.executorSummary(event.execId)
+      if (esummary.peakExecutorMetrics.compareAndUpdatePeakValues(event.executorMetrics)) {
+        update(esummary, now)
+      }
+    }
     // check if there is a new peak value for any of the executor level memory metrics,
     // while reading from the log. SparkListenerStageExecutorMetrics are only processed
     // when reading logs.
-    liveExecutors.get(executorMetrics.execId)
-      .orElse(deadExecutors.get(executorMetrics.execId)).map { exec =>
-      if (exec.peakExecutorMetrics.compareAndUpdatePeakValues(executorMetrics.executorMetrics)) {
+    liveExecutors.get(event.execId)
+      .orElse(deadExecutors.get(event.execId)).foreach { exec =>
+      if (exec.peakExecutorMetrics.compareAndUpdatePeakValues(event.executorMetrics)) {
         update(exec, now)
       }
     }

core/src/main/scala/org/apache/spark/status/AppStatusStore.scala

@@ -489,7 +489,8 @@ private[spark] class AppStatusStore(
       stage.accumulatorUpdates,
       Some(tasks),
       Some(executorSummary(stage.stageId, stage.attemptId)),
-      stage.killedTasksSummary)
+      stage.killedTasksSummary,
+      stage.peakExecutorMetrics)
   }
 
   def rdd(rddId: Int): v1.RDDStorageInfo = {

core/src/main/scala/org/apache/spark/status/LiveEntity.scala

@@ -325,6 +325,9 @@ private class LiveExecutorStageSummary(
 
   var metrics = createMetrics(default = 0L)
 
+  // peak values for executor level metrics
+  val peakExecutorMetrics = new ExecutorMetrics()
+
   override protected def doUpdate(): Any = {
     val info = new v1.ExecutorStageSummary(
       taskTime,
@@ -341,7 +344,8 @@ private class LiveExecutorStageSummary(
       metrics.shuffleWriteMetrics.recordsWritten,
       metrics.memoryBytesSpilled,
       metrics.diskBytesSpilled,
-      isBlacklisted)
+      isBlacklisted,
+      if (peakExecutorMetrics.isSet()) Some(peakExecutorMetrics) else None)
     new ExecutorStageSummaryWrapper(stageId, attemptId, executorId, info)
   }
 
@@ -380,6 +384,8 @@ private class LiveStage extends LiveEntity {
 
   var blackListedExecutors = new HashSet[String]()
 
+  val peakExecutorMetrics = new ExecutorMetrics()
+
   // Used for cleanup of tasks after they reach the configured limit. Not written to the store.
   @volatile var cleaning = false
   var savedTasks = new AtomicInteger(0)
@@ -429,7 +435,8 @@ private class LiveStage extends LiveEntity {
       newAccumulatorInfos(info.accumulables.values),
       None,
       None,
-      killedSummary)
+      killedSummary,
+      if (peakExecutorMetrics.isSet()) Some(peakExecutorMetrics) else None)
   }
 
   override protected def doUpdate(): Any = {

core/src/main/scala/org/apache/spark/status/api/v1/api.scala

@@ -76,7 +76,10 @@ class ExecutorStageSummary private[spark](
     val shuffleWriteRecords : Long,
     val memoryBytesSpilled : Long,
     val diskBytesSpilled : Long,
-    val isBlacklistedForStage: Boolean)
+    val isBlacklistedForStage: Boolean,
+    @JsonSerialize(using = classOf[ExecutorMetricsJsonSerializer])
+    @JsonDeserialize(using = classOf[ExecutorMetricsJsonDeserializer])
+    val peakMemoryMetrics: Option[ExecutorMetrics])
 
 class ExecutorSummary private[spark](
     val id: String,
@@ -236,7 +239,11 @@ class StageData private[spark](
     val accumulatorUpdates: Seq[AccumulableInfo],
     val tasks: Option[Map[Long, TaskData]],
     val executorSummary: Option[Map[String, ExecutorStageSummary]],
-    val killedTasksSummary: Map[String, Int])
+    val killedTasksSummary: Map[String, Int],
+
+    @JsonSerialize(using = classOf[ExecutorMetricsJsonSerializer])
+    @JsonDeserialize(using = classOf[ExecutorMetricsJsonDeserializer])
+    val peakExecutorMetrics: Option[ExecutorMetrics])
 
 class TaskData private[spark](
     val taskId: Long,

core/src/main/scala/org/apache/spark/ui/jobs/JobPage.scala

@@ -26,6 +26,7 @@ import scala.xml.{Node, NodeSeq, Unparsed, Utility}
 import org.apache.commons.lang3.StringEscapeUtils
 
 import org.apache.spark.JobExecutionStatus
+import org.apache.spark.executor.ExecutorMetrics
 import org.apache.spark.scheduler._
 import org.apache.spark.status.AppStatusStore
 import org.apache.spark.status.api.v1
@@ -216,7 +217,9 @@ private[ui] class JobPage(parent: JobsTab, store: AppStatusStore) extends WebUIP
           Nil,
           None,
           None,
-          Map())
+          Map(),
+          None
+        )
       }
     }