[SPARK-37580][CORE] Reset numFailures when one of task attempts succeeds

### What changes were proposed in this pull request?
When a task failed count reach the max threshold, abort check if another attempt succeed.

### Why are the changes needed?
In extreme situation, if  one task has failed 3 times(max failed threshold is 4 in default), and there is a retry task and speculative task both in running state, then one of these 2 task attempts succeed and to cancel another. But executor which task need to be cancelled lost(oom in our situcation), this task marked as failed, and TaskSetManager handle this failed task attempt, it has failed 4 times so abort this stage and cause job failed.

### Does this PR introduce _any_ user-facing change?
Yes, the meaning of `spark.task.maxFailures` has changed, from total count to continuous count of one particular task

### How was this patch tested?
Unit test.

Closes apache#34834 from wangshengjie123/fix_taskset_manager_abort_stage.

Lead-authored-by: wangshengjie3 <[email protected]>
Co-authored-by: wangshengjie <[email protected]>
Signed-off-by: yi.wu <[email protected]>

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

@@ -820,6 +820,7 @@ private[spark] class TaskSetManager(
         s"on ${info.host} (executor ${info.executorId}) ($tasksSuccessful/$numTasks)")
       // Mark successful and stop if all the tasks have succeeded.
       successful(index) = true
+      numFailures(index) = 0
       if (tasksSuccessful == numTasks) {
         isZombie = true
       }
@@ -843,6 +844,7 @@ private[spark] class TaskSetManager(
       if (!successful(index)) {
         tasksSuccessful += 1
         successful(index) = true
+        numFailures(index) = 0
         if (tasksSuccessful == numTasks) {
           isZombie = true
         }

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

@@ -2244,6 +2244,81 @@ class TaskSetManagerSuite
     // After 3s have elapsed now the task is marked as speculative task
     assert(sched.speculativeTasks.size == 1)
   }
+
+  test("SPARK-37580: Reset numFailures when one of task attempts succeeds") {
+    sc = new SparkContext("local", "test")
+    // Set the speculation multiplier to be 0 so speculative tasks are launched immediately
+    sc.conf.set(config.SPECULATION_MULTIPLIER, 0.0)
+    sc.conf.set(config.SPECULATION_QUANTILE, 0.6)
+    sc.conf.set(config.SPECULATION_ENABLED, true)
+
+    sched = new FakeTaskScheduler(sc, ("exec1", "host1"), ("exec2", "host2"), ("exec3", "host3"))
+    sched.backend = mock(classOf[SchedulerBackend])
+    val taskSet = FakeTask.createTaskSet(3)
+    val clock = new ManualClock()
+    val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES, clock = clock)
+
+    // Offer resources for 3 task to start
+    val tasks = new ArrayBuffer[TaskDescription]()
+    for ((k, v) <- List("exec1" -> "host1", "exec2" -> "host2", "exec3" -> "host3")) {
+      val taskOption = manager.resourceOffer(k, v, NO_PREF)._1
+      assert(taskOption.isDefined)
+      val task = taskOption.get
+      assert(task.executorId === k)
+      tasks += task
+    }
+    assert(sched.startedTasks.toSet === (0 until 3).toSet)
+
+    def runningTaskForIndex(index: Int): TaskDescription = {
+      tasks.find { task =>
+        task.index == index && !sched.endedTasks.contains(task.taskId)
+      }.getOrElse {
+        throw new RuntimeException(s"couldn't find index $index in " +
+          s"tasks: ${tasks.map { t => t.index -> t.taskId }} with endedTasks:" +
+          s" ${sched.endedTasks.keys}")
+      }
+    }
+    clock.advance(1)
+
+    // running task with index 1 fail 3 times (not enough to abort the stage)
+    (0 until 3).foreach { attempt =>
+      val task = runningTaskForIndex(1)
+      val endReason = ExceptionFailure("a", "b", Array(), "c", None)
+      manager.handleFailedTask(task.taskId, TaskState.FAILED, endReason)
+      sched.endedTasks(task.taskId) = endReason
+      assert(!manager.isZombie)
+      val nextTask = manager.resourceOffer(s"exec2", s"host2", NO_PREF)._1
+      assert(nextTask.isDefined, s"no offer for attempt $attempt of 1")
+      tasks += nextTask.get
+    }
+
+    val numFailuresField = classOf[TaskSetManager].getDeclaredField("numFailures")
+    numFailuresField.setAccessible(true)
+    val numFailures = numFailuresField.get(manager).asInstanceOf[Array[Int]]
+    // numFailures(1) should be 3
+    assert(numFailures(1) == 3)
+
+    // make task(TID 2) success to speculative other tasks
+    manager.handleSuccessfulTask(2, createTaskResult(2))
+
+    val originalTask = runningTaskForIndex(1)
+    clock.advance(1)
+    assert(manager.checkSpeculatableTasks(0))
+    assert(sched.speculativeTasks.toSet === Set(0, 1))
+
+    // make the speculative task(index 1) success
+    val speculativeTask = manager.resourceOffer("exec1", "host1", NO_PREF)._1
+    assert(speculativeTask.isDefined)
+    manager.handleSuccessfulTask(speculativeTask.get.taskId, createTaskResult(1))
+    // if task success, numFailures will be reset to 0
+    assert(numFailures(1) == 0)
+
+    // failed the originalTask(index 1) and check if the task manager is zombie
+    val failedReason = ExceptionFailure("a", "b", Array(), "c", None)
+    manager.handleFailedTask(originalTask.taskId, TaskState.FAILED, failedReason)
+    assert(!manager.isZombie)
+  }
+
 }
 
 class FakeLongTasks(stageId: Int, partitionId: Int) extends FakeTask(stageId, partitionId) {

docs/configuration.md

@@ -2461,9 +2461,10 @@ Apart from these, the following properties are also available, and may be useful
   <td><code>spark.task.maxFailures</code></td>
   <td>4</td>
   <td>
-    Number of failures of any particular task before giving up on the job.
+    Number of continuous failures of any particular task before giving up on the job.
     The total number of failures spread across different tasks will not cause the job
-    to fail; a particular task has to fail this number of attempts.
+    to fail; a particular task has to fail this number of attempts continuously.
+    If any attempt succeeds, the failure count for the task will be reset.
     Should be greater than or equal to 1. Number of allowed retries = this value - 1.
   </td>
   <td>0.8.0</td>