add stageIdToFinishedPartitions

Author: Ngone51

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

@@ -96,6 +96,9 @@ private[spark] class TaskSchedulerImpl(
   private[scheduler] val taskIdToTaskSetManager = new ConcurrentHashMap[Long, TaskSetManager]
   val taskIdToExecutorId = new HashMap[Long, String]
 
+  // Protected by `this`
+  private[scheduler] val stageIdToFinishedPartitions = new HashMap[Int, HashSet[Int]]
+
   @volatile private var hasReceivedTask = false
   @volatile private var hasLaunchedTask = false
   private val starvationTimer = new Timer(true)
@@ -208,7 +211,7 @@ private[spark] class TaskSchedulerImpl(
         taskSetsByStageIdAndAttempt.getOrElseUpdate(stage, new HashMap[Int, TaskSetManager])
       stageTaskSets(taskSet.stageAttemptId) = manager
       val conflictingTaskSet = stageTaskSets.exists { case (_, ts) =>
-        ts.taskSet != taskSet && !ts.isZombie
+        ts.taskSet != manager.taskSet && !ts.isZombie
       }
       if (conflictingTaskSet) {
         throw new IllegalStateException(s"more than one active taskSet for stage $stage:" +
@@ -238,7 +241,13 @@ private[spark] class TaskSchedulerImpl(
   private[scheduler] def createTaskSetManager(
       taskSet: TaskSet,
       maxTaskFailures: Int): TaskSetManager = {
-    new TaskSetManager(this, taskSet, maxTaskFailures, blacklistTrackerOpt)
+    val finishedPartitions =
+      stageIdToFinishedPartitions.getOrElseUpdate(taskSet.stageId, new HashSet[Int])
+    // filter the task which has been finished by previous attempts
+    val tasks = taskSet.tasks.filterNot{ t => finishedPartitions(t.partitionId) }
+    val ts = new TaskSet(
+      tasks, taskSet.stageId, taskSet.stageAttemptId, taskSet.priority, taskSet.properties)
+    new TaskSetManager(this, ts, maxTaskFailures, blacklistTrackerOpt)
   }
 
   override def cancelTasks(stageId: Int, interruptThread: Boolean): Unit = synchronized {
@@ -297,6 +306,7 @@ private[spark] class TaskSchedulerImpl(
       taskSetsForStage -= manager.taskSet.stageAttemptId
       if (taskSetsForStage.isEmpty) {
         taskSetsByStageIdAndAttempt -= manager.taskSet.stageId
+        stageIdToFinishedPartitions -= manager.taskSet.stageId
       }
     }
     manager.parent.removeSchedulable(manager)
@@ -837,17 +847,31 @@ private[spark] class TaskSchedulerImpl(
   }
 
   /**
-   * Marks the task has completed in all TaskSetManagers for the given stage.
+   * Marks the task has completed in all TaskSetManagers(active / zombie) for the given stage.
    *
    * After stage failure and retry, there may be multiple TaskSetManagers for the stage.
    * If an earlier attempt of a stage completes a task, we should ensure that the later attempts
    * do not also submit those same tasks.  That also means that a task completion from an earlier
    * attempt can lead to the entire stage getting marked as successful.
+   * And there's a situation that the active TaskSetManager corresponding to the stage may
+   * haven't been created at the time we call this method. And it is possible since the behaviour
+   * of calling on this method and creating active TaskSetManager is from two different threads,
+   * which are "task-result-getter" and "dag-scheduler-event-loop" separately. Consequently, under
+   * this situation, the active TaskSetManager which is created later could not learn about the
+   * finished partitions and keep on launching duplicate tasks, which may lead to job fail for some
+   * severe cases, see SPARK-25250 for details. So, to avoid the problem, we record the finished
+   * partitions for that stage here and exclude the already finished tasks when we creating active
+   * TaskSetManagers later by looking into stageIdToFinishedPartitions. Thus, active TaskSetManager
+   * could be always notified about the finished partitions whether it has been created or not at
+   * the time we call this method.
    */
   private[scheduler] def markPartitionCompletedInAllTaskSets(
       stageId: Int,
       partitionId: Int,
       taskInfo: TaskInfo) = {
+    val finishedPartitions =
+      stageIdToFinishedPartitions.getOrElseUpdate(stageId, new HashSet[Int])
+    finishedPartitions += partitionId
     taskSetsByStageIdAndAttempt.getOrElse(stageId, Map()).values.foreach { tsm =>
       tsm.markPartitionCompleted(partitionId, taskInfo)
     }

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

@@ -1206,6 +1206,67 @@ class TaskSchedulerImplSuite extends SparkFunSuite with LocalSparkContext with B
     }
   }
 
+  test("successful tasks from previous attempts could be learnt by later active taskset") {
+    val taskScheduler = setupSchedulerWithMockTaskSetBlacklist()
+    val valueSer = SparkEnv.get.serializer.newInstance()
+    val result = new DirectTaskResult[Int](valueSer.serialize(1), Seq())
+
+    // submit a taskset with 10 tasks to taskScheduler
+    val attempt0 = FakeTask.createTaskSet(10, stageId = 0, stageAttemptId = 0)
+    taskScheduler.submitTasks(attempt0)
+    // get the current active tsm
+    val tsm0 = taskScheduler.taskSetManagerForAttempt(0, 0).get
+    // offer sufficient resources
+    val offers0 = (0 until 10).map{ idx => WorkerOffer(s"exec-$idx", s"host-$idx", 1) }
+    taskScheduler.resourceOffers(offers0)
+    assert(tsm0.runningTasks === 10)
+    // fail task 0.0 and mark tsm0 as zombie
+    tsm0.handleFailedTask(tsm0.taskAttempts(0)(0).taskId, TaskState.FAILED,
+      FetchFailed(null, 0, 0, 0, "fetch failed"))
+    // the attempt0 is a zombie, but the tasks are still running (this could be true even if
+    // we actively killed those tasks, as killing is best-effort)
+    assert(tsm0.isZombie)
+    assert(tsm0.runningTasks === 9)
+
+
+    // success task 1.0 , finish partition 1. But now,
+    // no active tsm exists in TaskScheduler for stage0.
+    tsm0.handleSuccessfulTask(tsm0.taskAttempts(1)(0).taskId, result)
+    assert(tsm0.runningTasks === 8)
+    assert(taskScheduler.stageIdToFinishedPartitions(0).contains(1))
+
+    // submit a new taskset with 10 tasks after someone previous task attempt succeed
+    val attempt1 = FakeTask.createTaskSet(10, stageId = 0, stageAttemptId = 1)
+    taskScheduler.submitTasks(attempt1)
+    // get the current active tsm
+    val tsm1 = taskScheduler.taskSetManagerForAttempt(0, 1).get
+    // tsm1 learns about the finished partition 1 during constructing, so it only need
+    // to execute other 9 tasks
+    assert(tsm1.taskSet.tasks.length == 9)
+    // offer one resource
+    val offers1 = (10 until 11).map{ idx => WorkerOffer(s"exec-$idx", s"host-$idx", 1) }
+    taskScheduler.resourceOffers(offers1)
+    assert(tsm1.runningTasks === 1)
+    // success task 0.0 in tsm1 and finish partition 0
+    tsm1.handleSuccessfulTask(tsm1.taskAttempts(0)(0).taskId, result)
+    assert(taskScheduler.stageIdToFinishedPartitions(0).contains(0))
+
+
+    val runningTasks = tsm0.taskSet.tasks.filterNot{ t =>
+      taskScheduler.stageIdToFinishedPartitions(0).contains(t.partitionId)
+    }
+    // finish tsm1 by previous task attempts from tsm0, this remains same behavior with SPARK-23433
+    runningTasks.foreach{ t =>
+      val attempt = tsm0.taskAttempts(tsm0.partitionToIndex(t.partitionId)).head
+      tsm0.handleSuccessfulTask(attempt.taskId, result)
+    }
+
+    assert(taskScheduler.taskSetManagerForAttempt(0, 0).isEmpty)
+    assert(taskScheduler.taskSetManagerForAttempt(0, 1).isEmpty)
+    assert(taskScheduler.stageIdToFinishedPartitions.isEmpty)
+
+  }
+
   test("don't schedule for a barrier taskSet if available slots are less than pending tasks") {
     val taskCpus = 2
     val taskScheduler = setupScheduler(config.CPUS_PER_TASK.key -> taskCpus.toString)