[CARMEL-7340][CARMEL-3697] Task Summary Support in TaskScheduler (apache#100)

Author: wangyum

core/src/main/scala/org/apache/spark/scheduler/AnalyticsTaskSetManager.scala

@@ -100,8 +100,6 @@ private[spark] class AnalyticsTaskSetManager(
 
   def userInfo(): Option[UserInfo] = _userInfo
 
-  def successfulTasks: Int = tasksSuccessful
-
   // Add all our tasks to the pending lists. We do this in reverse order
   // of task index so that tasks with low indices get launched first.
   addPendingTasks()

core/src/main/scala/org/apache/spark/scheduler/TaskScheduler.scala

@@ -125,4 +125,5 @@ private[spark] trait TaskScheduler {
    */
   def applicationAttemptId(): Option[String]
 
+  def taskSummary(): (Int, Int, Int)
 }

core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala

@@ -1258,6 +1258,20 @@ private[spark] class TaskSchedulerImpl(
       manager
     }
   }
+
+  override def taskSummary(): (Int, Int, Int) = {
+    var totalTasks = 0
+    var runningTasks = 0
+    var successfulTasks = 0
+    for (taskSet <- rootPool.getTaskSetQueue) {
+      if (!taskSet.isZombie) {
+        totalTasks += taskSet.numTasks
+        runningTasks += taskSet.runningTasks
+        successfulTasks += taskSet.successfulTasks
+      }
+    }
+    (totalTasks, runningTasks, successfulTasks)
+  }
 }
 
 

core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala

@@ -172,6 +172,8 @@ private[spark] class TaskSetManager(
 
   override def runningTasks: Int = runningTasksSet.size
 
+  def successfulTasks: Int = tasksSuccessful
+
   def someAttemptSucceeded(tid: Long): Boolean = {
     successful(taskInfos(tid).index)
   }

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

@@ -593,16 +593,21 @@ private[spark] class AppStatusStore(
     constructTaskDataList(taskDataWrapperSeq)
   }
 
-  def resourceUsage(): v1.ClusterResourceSummary = {
+  def resourceUsage(sc: Option[SparkContext]): v1.ClusterResourceSummary = {
     val executors = store.view(classOf[ExecutorSummaryWrapper]).asScala
       .filter(d => d.info.isActive)
       .map(_.info.totalCores)
       .sum
-    val tasks = store.view(classOf[StageDataWrapper]).reverse().asScala
-      .filter(d => d.info.status.equals(v1.StageStatus.ACTIVE))
-    val numTasks = tasks.map(_.info.numTasks).sum
-    val numActiveTasks = tasks.map(_.info.numActiveTasks).sum
-    val numCompleteTasks = tasks.map(_.info.numCompleteTasks).sum
+    val (numTasks, numActiveTasks, numCompleteTasks) = if (sc.isDefined) {
+      sc.get.taskScheduler.taskSummary()
+    } else {
+      val tasks = store.view(classOf[StageDataWrapper]).reverse().asScala
+        .filter(d => d.info.status.equals(v1.StageStatus.ACTIVE))
+      val numTasks = tasks.map(_.info.numTasks).sum
+      val numActiveTasks = tasks.map(_.info.numActiveTasks).sum
+      val numCompleteTasks = tasks.map(_.info.numCompleteTasks).sum
+      (numTasks, numActiveTasks, numCompleteTasks)
+    }
 
     new v1.ClusterResourceSummary(executors, numTasks, numActiveTasks, numCompleteTasks)
   }

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

@@ -174,8 +174,7 @@ private[v1] class AbstractApplicationResource extends BaseAppResource {
 
   @GET
   @Path("cluster/usage")
-  def executorSimpleList(): ClusterResourceSummary = withUI(_.store.resourceUsage())
-
+  def executorSimpleList(): ClusterResourceSummary = withUI(ui => ui.store.resourceUsage(ui.sc))
 }
 
 private[v1] class OneApplicationResource extends AbstractApplicationResource {

core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala

@@ -207,6 +207,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti
     override def executorLost(executorId: String, reason: ExecutorLossReason): Unit = {}
     override def workerRemoved(workerId: String, host: String, message: String): Unit = {}
     override def applicationAttemptId(): Option[String] = None
+    override def taskSummary(): (Int, Int, Int) = (0, 0, 0)
     override def executorDecommission(
       executorId: String,
       decommissionInfo: ExecutorDecommissionInfo): Unit = {
@@ -890,6 +891,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti
       override def executorLost(executorId: String, reason: ExecutorLossReason): Unit = {}
       override def workerRemoved(workerId: String, host: String, message: String): Unit = {}
       override def applicationAttemptId(): Option[String] = None
+      override def taskSummary(): (Int, Int, Int) = (0, 0, 0)
       override def executorDecommission(
         executorId: String,
         decommissionInfo: ExecutorDecommissionInfo): Unit = {}

core/src/test/scala/org/apache/spark/scheduler/ExternalClusterManagerSuite.scala

@@ -93,6 +93,7 @@ private class DummyTaskScheduler extends TaskScheduler {
   override def executorLost(executorId: String, reason: ExecutorLossReason): Unit = {}
   override def workerRemoved(workerId: String, host: String, message: String): Unit = {}
   override def applicationAttemptId(): Option[String] = None
+  override def taskSummary(): (Int, Int, Int) = (0, 0, 0)
   def executorHeartbeatReceived(
       execId: String,
       accumUpdates: Array[(Long, Seq[AccumulatorV2[_, _]])],

core/src/test/scala/org/apache/spark/storage/BlockManagerDecommissionIntegrationSuite.scala

@@ -330,10 +330,12 @@ class BlockManagerDecommissionIntegrationSuite extends SparkFunSuite with LocalS
     val execIdToBlocksMapping = storageStatus.map(
       status => (status.blockManagerId.executorId, status.blocks)).toMap
     // No cached blocks should be present on executor which was decommissioned
-    assert(
-      !execIdToBlocksMapping.contains(execToDecommission) ||
-      execIdToBlocksMapping(execToDecommission).keys.filter(_.isRDD).toSeq === Seq(),
-      "Cache blocks should be migrated")
+    eventually(timeout(2.seconds), interval(100.milliseconds)) {
+      assert(
+        !execIdToBlocksMapping.contains(execToDecommission) ||
+          execIdToBlocksMapping(execToDecommission).keys.filter(_.isRDD).toSeq === Seq(),
+        "Cache blocks should be migrated")
+    }
     if (persist) {
       // There should still be all the RDD blocks cached
       assert(execIdToBlocksMapping.values.flatMap(_.keys).count(_.isRDD) === numParts)