Compute numTasks from job start stage infos.

Author: JoshRosen

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

@@ -85,6 +85,12 @@ class JobProgressListener(conf: SparkConf) extends SparkListener with Logging {
     val jobData: JobUIData =
       new JobUIData(jobStart.jobId, Some(System.currentTimeMillis), None, jobStart.stageIds,
         jobGroup, JobExecutionStatus.RUNNING)
+    // Compute (a potential underestimate of) the number of tasks that will be run by this job:
+    jobData.numTasks = {
+      val allStages = jobStart.stageInfos
+      val missingStages = allStages.filter(_.completionTime.isEmpty)
+      missingStages.map(_.numTasks).sum
+    }
     jobIdToData(jobStart.jobId) = jobData
     activeJobs(jobStart.jobId) = jobData
     for (stageId <- jobStart.stageIds) {
@@ -189,7 +195,6 @@ class JobProgressListener(conf: SparkConf) extends SparkListener with Logging {
       jobId <- activeJobsDependentOnStage;
       jobData <- jobIdToData.get(jobId)
     ) {
-      jobData.numTasks += stage.numTasks
       jobData.numActiveStages += 1
     }
   }
@@ -245,8 +250,6 @@ class JobProgressListener(conf: SparkConf) extends SparkListener with Logging {
       execSummary.taskTime += info.duration
       stageData.numActiveTasks -= 1
 
-      val isRecomputation = stageData.completedIndices.contains(info.index)
-
       val (errorMessage, metrics): (Option[String], Option[TaskMetrics]) =
         taskEnd.reason match {
           case org.apache.spark.Success =>
@@ -279,9 +282,7 @@ class JobProgressListener(conf: SparkConf) extends SparkListener with Logging {
         jobData.numActiveTasks -= 1
         taskEnd.reason match {
           case Success =>
-            if (!isRecomputation) {
-              jobData.numCompletedTasks += 1
-            }
+            jobData.numCompletedTasks += 1
           case _ =>
             jobData.numFailedTasks += 1
         }

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

@@ -46,6 +46,8 @@ private[jobs] object UIData {
     var jobGroup: Option[String] = None,
     var status: JobExecutionStatus = JobExecutionStatus.UNKNOWN,
     /* Tasks */
+    // `numTasks` is a potential underestimate of the true number of tasks that this job will run
+    // see https://github.com/apache/spark/pull/3009 for an extensive discussion of this
     var numTasks: Int = 0,
     var numActiveTasks: Int = 0,
     var numCompletedTasks: Int = 0,

core/src/test/scala/org/apache/spark/ui/UISeleniumSuite.scala

@@ -188,6 +188,12 @@ class UISeleniumSuite extends FunSuite with WebBrowser with Matchers {
       eventually(timeout(5 seconds), interval(50 milliseconds)) {
         go to (sc.ui.get.appUIAddress.stripSuffix("/") + "/jobs")
         find(cssSelector(".stage-progress-cell")).get.text should be ("2/2 (1 failed)")
+        // Ideally, the following test would pass, but currently we over-count completed tasks
+        // if task recomputations occur:
+        // find(cssSelector(".progress-cell .progress")).get.text should be ("2/2 (1 failed)")
+        // Instead, we guarantee that the total number of tasks is always correct, while the number
+        // of completed tasks may be higher:
+        find(cssSelector(".progress-cell .progress")).get.text should be ("3/2 (1 failed)")
       }
     }
   }