[SPARK-39955][CORE] Improve LaunchTask process to avoid Stage failures caused by fail-to-send LaunchTask messages

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

@@ -93,6 +93,8 @@ class TaskInfo(
 
   var killed = false
 
+  var launching = true
+
   private[spark] def markGettingResult(time: Long): Unit = {
     gettingResultTime = time
   }
@@ -108,6 +110,10 @@ class TaskInfo(
     }
   }
 
+  private[spark] def launchSucceeded(): Unit = {
+    launching = false
+  }
+
   def gettingResult: Boolean = gettingResultTime != 0
 
   def finished: Boolean = finishTime != 0

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

@@ -825,6 +825,9 @@ private[spark] class TaskSchedulerImpl(
                 taskResultGetter.enqueueFailedTask(taskSet, tid, state, serializedData)
               }
             }
+            if (state == TaskState.RUNNING) {
+              taskSet.taskInfos(tid).launchSucceeded()
+            }
           case None =>
             logError(
               ("Ignoring update with state %s for TID %s because its task set is gone (this is " +

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

@@ -1068,10 +1068,13 @@ private[spark] class TaskSetManager(
     }
     for ((tid, info) <- taskInfos if info.running && info.executorId == execId) {
       val exitCausedByApp: Boolean = reason match {
-        case exited: ExecutorExited => exited.exitCausedByApp
+        case ExecutorExited(_, false, _) => false
         case ExecutorKilled | ExecutorDecommission(_) => false
         case ExecutorProcessLost(_, _, false) => false
-        case _ => true
+        // If the task is launching, this indicates that Driver has sent LaunchTask to Executor,
+        // but Executor has not sent StatusUpdate(TaskState.RUNNING) to Driver. Hence, we assume
+        // that the task is not running, and it is NetworkFailure rather than TaskFailure.
+        case _ => !info.launching
       }
       handleFailedTask(tid, TaskState.FAILED, ExecutorLostFailure(info.executorId, exitCausedByApp,
         Some(reason.toString)))

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

@@ -2085,6 +2085,53 @@ class TaskSchedulerImplSuite extends SparkFunSuite with LocalSparkContext
     assert(!taskSetManager.successful(taskDescriptions(0).index))
   }
 
+  Seq(true, false).foreach { hasLaunched =>
+    val testName = if (hasLaunched) {
+      "executor lost could fail task set if task is running"
+    } else {
+      "executor lost should not fail task set if task is launching"
+    }
+    test(s"SPARK-39955: $testName") {
+      val taskCpus = 2
+      val taskScheduler = setupSchedulerWithMaster(
+        s"local[$taskCpus]",
+        config.TASK_MAX_FAILURES.key -> "1")
+      taskScheduler.initialize(new FakeSchedulerBackend)
+      // Need to initialize a DAGScheduler for the taskScheduler to use for callbacks.
+      new DAGScheduler(sc, taskScheduler) {
+        override def taskStarted(task: Task[_], taskInfo: TaskInfo): Unit = {}
+        override def executorAdded(execId: String, host: String): Unit = {}
+      }
+
+      val workerOffer = IndexedSeq(
+        WorkerOffer("executor0", "host0", 1))
+      val taskSet = FakeTask.createTaskSet(1)
+      // submit tasks, offer resources, task gets scheduled
+      taskScheduler.submitTasks(taskSet)
+      var tsm: Option[TaskSetManager] = None
+      eventually(timeout(10.seconds)) {
+        tsm = taskScheduler.taskSetManagerForAttempt(taskSet.stageId, taskSet.stageAttemptId)
+        assert(tsm.isDefined && !tsm.get.isZombie)
+      }
+      val taskDescriptions = taskScheduler.resourceOffers(workerOffer)
+      assert(1 === taskDescriptions.length)
+      assert(taskScheduler.runningTasksByExecutors("executor0") === 1)
+      if (hasLaunched) {
+        taskScheduler.statusUpdate(
+          0,
+          TaskState.RUNNING,
+          ByteBuffer.allocate(0))
+        eventually(timeout(10.seconds)) {
+          assert(!tsm.get.taskInfos(0).launching)
+        }
+      }
+      taskScheduler.executorLost("executor0", ExecutorProcessLost())
+      eventually(timeout(10.seconds)) {
+        assert(tsm.get.isZombie === hasLaunched)
+      }
+    }
+  }
+
   /**
    * Used by tests to simulate a task failure. This calls the failure handler explicitly, to ensure
    * that all the state is updated when this method returns. Otherwise, there's no way to know when

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

@@ -639,6 +639,11 @@ class TaskSetManagerSuite
       ExecutorExited(143, false, "Terminated for reason unrelated to running tasks"))
     assert(!sched.taskSetsFailed.contains(taskSet.id))
     assert(manager.resourceOffer("execC", "host2", ANY)._1.isDefined)
+
+    // Driver receives StatusUpdate(RUNNING) from Executors
+    for ((tid, info) <- manager.taskInfos if info.running) {
+      manager.taskInfos(tid).launchSucceeded()
+    }
     sched.removeExecutor("execC")
     manager.executorLost(
       "execC", "host2", ExecutorExited(1, true, "Terminated due to issue with running tasks"))
@@ -2202,6 +2207,7 @@ class TaskSetManagerSuite
 
     val (taskId0, index0, exec0) = (task0.taskId, task0.index, task0.executorId)
     val (taskId1, index1, exec1) = (task1.taskId, task1.index, task1.executorId)
+
     // set up two running tasks
     assert(manager.taskInfos(taskId0).running)
     assert(manager.taskInfos(taskId1).running)
@@ -2211,6 +2217,12 @@ class TaskSetManagerSuite
     assert(manager.invokePrivate(numFailures())(index0) === 0)
     assert(manager.invokePrivate(numFailures())(index1) === 0)
 
+    sched.asInstanceOf[TaskSchedulerImpl]
+      .statusUpdate(taskId1, TaskState.RUNNING, ByteBuffer.allocate(0))
+    eventually(timeout(10.seconds), interval(100.milliseconds)) {
+      assert(manager.taskInfos(taskId0).running)
+      assert(manager.taskInfos(taskId1).running && !manager.taskInfos(taskId1).launching)
+    }
     // let exec1 count task failures but exec0 doesn't
     backend.executorsPendingToRemove(exec0) = true
     backend.executorsPendingToRemove(exec1) = false