[SPARK-59418][CORE] Prefer idle executors and isolate repeated OOM retries

core/src/main/scala/org/apache/spark/ExecutorAllocationManager.scala

@@ -108,7 +108,8 @@ private[spark] class ExecutorAllocationManager(
     cleaner: Option[ContextCleaner] = None,
     clock: Clock = new SystemClock(),
     resourceProfileManager: ResourceProfileManager,
-    reliableShuffleStorage: Boolean)
+    reliableShuffleStorage: Boolean,
+    oomRetryReservationInfo: () => Option[OomRetryReservationInfo] = () => None)
   extends Logging {
 
   allocationManager =>
@@ -135,6 +136,8 @@ private[spark] class ExecutorAllocationManager(
 
   private val decommissionEnabled = conf.get(DECOMMISSION_ENABLED)
 
+  private val oomRetryEnabled = conf.get(SCHEDULER_OOM_RETRY_ENABLED)
+
   private val defaultProfileId = resourceProfileManager.defaultResourceProfile.id
 
   validateSettings()
@@ -413,7 +416,7 @@ private[spark] class ExecutorAllocationManager(
    * The maximum number of executors, for the ResourceProfile id passed in, that we would need
    * under the current load to satisfy all running and pending tasks, rounded up.
    */
-  private[spark] def maxNumExecutorsNeededPerResourceProfile(rpId: Int): Int = {
+  private[spark] def maxNumExecutorsNeededPerResourceProfile(rpId: Int): Int = synchronized {
     val pendingTask = listener.pendingTasksPerResourceProfile(rpId)
     val pendingSpeculative = listener.pendingSpeculativeTasksPerResourceProfile(rpId)
     val unschedulableTaskSets = listener.pendingUnschedulableTaskSetsPerResourceProfile(rpId)
@@ -423,16 +426,28 @@ private[spark] class ExecutorAllocationManager(
     val tasksPerExecutor = rp.maxTasksPerExecutor(conf)
     logDebug(s"max needed for rpId: $rpId numpending: $numRunningOrPendingTasks," +
       s" tasksperexecutor: $tasksPerExecutor")
-    val maxNeeded = math.ceil(numRunningOrPendingTasks * executorAllocationRatio /
-      tasksPerExecutor).toInt
-
-    val maxNeededWithSpeculationLocalityOffset =
-      if (tasksPerExecutor > 1 && maxNeeded == 1 && pendingSpeculative > 0) {
-      // If we have pending speculative tasks and only need a single executor, allocate one more
-      // to satisfy the locality requirements of speculation
-      maxNeeded + 1
+    def executorsNeeded(tasks: Int, speculationNeedsSeparation: Boolean): Int = {
+      val needed = math.ceil(tasks * executorAllocationRatio / tasksPerExecutor).toInt
+      if (tasksPerExecutor > 1 && needed == 1 && speculationNeedsSeparation) {
+        // A speculative copy needs an executor other than the one running its original attempt.
+        needed + 1
+      } else {
+        needed
+      }
+    }
+
+    val baseline = executorsNeeded(numRunningOrPendingTasks, pendingSpeculative > 0)
+    val maxNeeded = if (oomRetryEnabled && tasksPerExecutor > 1) {
+      oomRetryReservationInfo().filter(_.resourceProfileId == rpId).map { reservation =>
+        // The reservation occupies one whole executor, even at a reduced allocation ratio.
+        // Size the remaining work separately, including its speculative-locality requirement.
+        val (coveredTasks, speculationNeedsSeparation) =
+          listener.oomRetryAllocation(reservation)
+        val ordinaryTasks = math.max(0, numRunningOrPendingTasks - coveredTasks)
+        math.max(baseline, 1 + executorsNeeded(ordinaryTasks, speculationNeedsSeparation))
+      }.getOrElse(baseline)
     } else {
-      maxNeeded
+      baseline
     }
 
     if (unschedulableTaskSets > 0) {
@@ -441,10 +456,10 @@ private[spark] class ExecutorAllocationManager(
       // the max needed which we would normally get.
       val maxNeededForUnschedulables = math.ceil(unschedulableTaskSets * executorAllocationRatio /
         tasksPerExecutor).toInt
-      math.max(maxNeededWithSpeculationLocalityOffset,
+      math.max(maxNeeded,
         executorMonitor.executorCountWithResourceProfile(rpId) + maxNeededForUnschedulables)
     } else {
-      maxNeededWithSpeculationLocalityOffset
+      maxNeeded
     }
   }
 
@@ -791,6 +806,9 @@ private[spark] class ExecutorAllocationManager(
     // Number of running tasks per stageAttempt including speculative tasks.
     // Should be 0 when no stages are active.
     private val stageAttemptToNumRunningTask = new mutable.HashMap[StageAttempt, Int]
+    // Used only for OOM recovery. Keep coverage in the same event domain as allocation demand.
+    private val executorToRunningTasks =
+      new mutable.HashMap[String, mutable.HashMap[Long, (StageAttempt, Int, Boolean)]]
     private val stageAttemptToTaskIndices = new mutable.HashMap[StageAttempt, mutable.HashSet[Int]]
     // Map from each stageAttempt to a set of running speculative task indexes
     // TODO(SPARK-41192): We simply need an Int for this.
@@ -901,6 +919,11 @@ private[spark] class ExecutorAllocationManager(
       val stageAttempt = StageAttempt(stageId, stageAttemptId)
       val taskIndex = taskStart.taskInfo.index
       allocationManager.synchronized {
+        if (oomRetryEnabled) {
+          executorToRunningTasks.getOrElseUpdate(taskStart.taskInfo.executorId,
+            new mutable.HashMap)(taskStart.taskInfo.taskId) =
+            (stageAttempt, taskIndex, taskStart.taskInfo.speculative)
+        }
         stageAttemptToNumRunningTask(stageAttempt) =
           stageAttemptToNumRunningTask.getOrElse(stageAttempt, 0) + 1
         // If this is the last pending task, mark the scheduler queue as empty
@@ -925,7 +948,15 @@ private[spark] class ExecutorAllocationManager(
       val stageAttempt = StageAttempt(stageId, stageAttemptId)
       val taskIndex = taskEnd.taskInfo.index
       allocationManager.synchronized {
-        if (stageAttemptToNumRunningTask.contains(stageAttempt)) {
+        executorToRunningTasks.get(taskEnd.taskInfo.executorId).foreach { runningTasks =>
+          runningTasks.remove(taskEnd.taskInfo.taskId)
+          if (runningTasks.isEmpty) {
+            executorToRunningTasks.remove(taskEnd.taskInfo.executorId)
+          }
+        }
+        // Resubmitted reports lost output from a task that already finished, not another exit.
+        if (taskEnd.reason != Resubmitted &&
+            stageAttemptToNumRunningTask.contains(stageAttempt)) {
           stageAttemptToNumRunningTask(stageAttempt) -= 1
           if (stageAttemptToNumRunningTask(stageAttempt) == 0) {
             stageAttemptToNumRunningTask -= stageAttempt
@@ -943,14 +974,25 @@ private[spark] class ExecutorAllocationManager(
             stageAttemptToPendingSpeculativeTasks.get(stageAttempt).foreach(_.remove(taskIndex))
           case _: TaskKilled =>
           case _ =>
+            val isOom = taskEnd.reason match {
+              case e: ExceptionFailure => e.isOutOfMemoryError
+              case e: ExecutorLostFailure => e.exitCausedByApp && e.isOutOfMemoryError
+              case _ => false
+            }
+            if (oomRetryEnabled && isOom) {
+              // TaskSetManager revokes speculation for OOM-affected partitions. Do not retain
+              // executor demand for a speculative attempt that can no longer be scheduled.
+              stageAttemptToPendingSpeculativeTasks.get(stageAttempt).foreach(_.remove(taskIndex))
+            }
             if (!hasPendingTasks) {
               // If the task failed (not intentionally killed), we expect it to be resubmitted
               // later. To ensure we have enough resources to run the resubmitted task, we need to
               // mark the scheduler as backlogged again if it's not already marked as such
               // (SPARK-8366)
               allocationManager.onSchedulerBacklogged()
             }
-            if (!taskEnd.taskInfo.speculative) {
+            // Lost shuffle output is retried as a regular task even if a speculative copy won.
+            if (!taskEnd.taskInfo.speculative || taskEnd.reason == Resubmitted) {
               // If a non-speculative task is intentionally killed, it means the speculative task
               // has succeeded, and no further task of this task index will be resubmitted. In this
               // case, the task index is completed and we shouldn't remove it from
@@ -1080,6 +1122,32 @@ private[spark] class ExecutorAllocationManager(
       }.sum
     }
 
+    /** Count only reserved tasks that also contribute to this listener's demand. */
+    def oomRetryAllocation(reservation: OomRetryReservationInfo): (Int, Boolean) = {
+      val attempts = resourceProfileIdToStageAttempt
+        .getOrElse(reservation.resourceProfileId, mutable.Set.empty[StageAttempt])
+      val coveredRunning = executorToRunningTasks.get(reservation.executorId).iterator
+        .flatMap(_.valuesIterator).filter { case (attempt, _, _) => attempts.contains(attempt) }
+        .toSeq
+      val retryAttempt = StageAttempt(reservation.stageId, reservation.stageAttemptId)
+      val coveredPending = if (attempts.contains(retryAttempt) &&
+          stageAttemptToNumTasks.contains(retryAttempt) &&
+          !stageAttemptToTaskIndices.get(retryAttempt).exists(_.contains(reservation.taskIndex))) {
+        1
+      } else {
+        0
+      }
+      val coveredOriginals = coveredRunning.collect {
+        case (attempt, index, false) => (attempt, index)
+      }.toSet
+      val speculationNeedsSeparation = attempts.exists { attempt =>
+        stageAttemptToPendingSpeculativeTasks.get(attempt).exists { indices =>
+          indices.exists(index => !coveredOriginals.contains((attempt, index)))
+        }
+      }
+      (coveredRunning.size + coveredPending, speculationNeedsSeparation)
+    }
+
     /**
      * Update the Executor placement hints (the number of tasks with locality preferences,
      * a map where each pair is a node and the number of tasks that would like to be scheduled

core/src/main/scala/org/apache/spark/SparkContext.scala

@@ -726,10 +726,12 @@ class SparkContext(config: SparkConf) extends Logging {
       if (dynamicAllocationEnabled) {
         schedulerBackend match {
           case b: ExecutorAllocationClient =>
+            val taskScheduler = _taskScheduler
             Some(new ExecutorAllocationManager(
               schedulerBackend.asInstanceOf[ExecutorAllocationClient], listenerBus, _conf,
               cleaner = cleaner, resourceProfileManager = resourceProfileManager,
-              reliableShuffleStorage = _shuffleDriverComponents.supportsReliableStorage()))
+              reliableShuffleStorage = _shuffleDriverComponents.supportsReliableStorage(),
+              oomRetryReservationInfo = () => taskScheduler.oomRetryReservationInfo))
           case _ =>
             None
         }

core/src/main/scala/org/apache/spark/TaskEndReason.scala

@@ -21,6 +21,7 @@ import java.io.{ObjectInputStream, ObjectOutputStream}
 
 import org.apache.spark.annotation.DeveloperApi
 import org.apache.spark.internal.Logging
+import org.apache.spark.memory.SparkOutOfMemoryError
 import org.apache.spark.scheduler.AccumulableInfo
 import org.apache.spark.storage.BlockManagerId
 import org.apache.spark.util.{AccumulatorV2, Utils}
@@ -135,6 +136,11 @@ case class ExceptionFailure(
     private[spark] var metricPeaks: Seq[Long] = Seq.empty)
   extends TaskFailedReason {
 
+  // Keep this independent of exceptionWrapper, which may be dropped during serialization.
+  private[spark] var isOutOfMemoryError: Boolean =
+    className == classOf[OutOfMemoryError].getName ||
+      className == classOf[SparkOutOfMemoryError].getName
+
   /**
    * `preserveCause` is used to keep the exception itself so it is available to the
    * driver. This may be set to `false` in the event that the exception is not in fact
@@ -146,6 +152,14 @@ case class ExceptionFailure(
       preserveCause: Boolean) = {
     this(e.getClass.getName, e.getMessage, e.getStackTrace, Utils.exceptionString(e),
       if (preserveCause) Some(new ThrowableSerializationWrapper(e)) else None, accumUpdates)
+    // Bound the traversal because user exceptions may contain cycles in their cause chains.
+    var cause = e
+    var depth = 0
+    while (cause != null && !isOutOfMemoryError && depth < 20) {
+      isOutOfMemoryError = cause.isInstanceOf[OutOfMemoryError]
+      cause = cause.getCause
+      depth += 1
+    }
   }
 
   private[spark] def this(e: Throwable, accumUpdates: Seq[AccumulableInfo]) = {
@@ -262,6 +276,8 @@ case class ExecutorLostFailure(
     execId: String,
     exitCausedByApp: Boolean = true,
     reason: Option[String]) extends TaskFailedReason {
+  private[spark] var isOutOfMemoryError: Boolean = false
+
   override def toErrorString: String = {
     val exitBehavior = if (exitCausedByApp) {
       "caused by one of the running tasks"

core/src/main/scala/org/apache/spark/internal/config/package.scala

@@ -2639,6 +2639,29 @@ package object config {
       .intConf
       .createWithDefault(5)
 
+  private[spark] val SCHEDULER_OOM_RETRY_ENABLED =
+    ConfigBuilder("spark.scheduler.oomRetry.enabled")
+      .doc("Prefer idle executors for tasks retried after an out-of-memory failure. After two " +
+        "OOM failures of a task, temporarily reserve an executor so its retry can run alone. " +
+        "At most one executor per application is reserved, and existing tasks are allowed to " +
+        "finish before the retry starts. OOM retries may ignore preferred locations, but still " +
+        "respect exclusions and resource requirements. Task CPUs and executor memory are " +
+        "unchanged. Barrier and pipelined tasks are excluded and OOM retries are not speculated.")
+      .version("5.0.0")
+      .booleanConf
+      .createWithDefault(false)
+
+  private[spark] val SCHEDULER_OOM_RETRY_ISOLATION_TIMEOUT =
+    ConfigBuilder("spark.scheduler.oomRetry.isolationTimeout")
+      .doc("Maximum wait for an isolated retry after an OOM failure when " +
+        "spark.scheduler.oomRetry.enabled is true. After this time the pending retry falls " +
+        "back to ordinary placement; another OOM failure starts a new wait. This does not " +
+        "limit the running time of a retry that has already started in isolation.")
+      .version("5.0.0")
+      .timeConf(TimeUnit.MILLISECONDS)
+      .checkValue(_ > 0, "OOM retry isolation timeout must be positive")
+      .createWithDefaultString("60s")
+
   private[spark] val SCHEDULER_REVIVE_INTERVAL =
     ConfigBuilder("spark.scheduler.revive.interval")
       .version("0.8.1")

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

@@ -18,6 +18,7 @@
 package org.apache.spark.scheduler
 
 import org.apache.spark.executor.ExecutorExitCode
+import org.apache.spark.util.SparkExitCode
 
 /**
  * Represents an explanation for an executor or whole process failing or exiting.
@@ -29,7 +30,10 @@ class ExecutorLossReason(val message: String) extends Serializable {
 
 private[spark]
 case class ExecutorExited(exitCode: Int, exitCausedByApp: Boolean, reason: String)
-  extends ExecutorLossReason(reason)
+  extends ExecutorLossReason(reason) {
+  private[spark] var isOutOfMemoryError: Boolean =
+    exitCausedByApp && exitCode == SparkExitCode.OOM
+}
 
 private[spark] object ExecutorExited {
   def apply(exitCode: Int, exitCausedByApp: Boolean): ExecutorExited = {

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

@@ -74,6 +74,23 @@ private[spark] class ExecutorResourcesAmounts(
     }
   }
 
+  /** Count tasks that fit the current resource amounts, without consuming the live offer. */
+  def availableTaskSlots(taskSetProf: ResourceProfile, maxTasks: Int): Int = {
+    if (taskSetProf.getCustomTaskResources().isEmpty) return maxTasks
+    val remaining = new ExecutorResourcesAmounts(
+      internalResources.map { case (name, amounts) => name -> amounts.toMap })
+    var slots = 0
+    while (slots < maxTasks) {
+      remaining.assignAddressesCustomResources(taskSetProf) match {
+        case Some(assigned) =>
+          remaining.acquire(assigned)
+          slots += 1
+        case None => return slots
+      }
+    }
+    slots
+  }
+
   /**
    * Acquire the resource.
    * @param assignedResource the assigned resource information

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

@@ -24,6 +24,13 @@ import org.apache.spark.scheduler.SchedulingMode.SchedulingMode
 import org.apache.spark.storage.BlockManagerId
 import org.apache.spark.util.AccumulatorV2
 
+private[spark] case class OomRetryReservationInfo(
+    resourceProfileId: Int,
+    executorId: String,
+    stageId: Int,
+    stageAttemptId: Int,
+    taskIndex: Int)
+
 /**
  * Low-level task scheduler interface, currently implemented exclusively by
  * [[org.apache.spark.scheduler.TaskSchedulerImpl]].
@@ -76,6 +83,13 @@ private[spark] trait TaskScheduler {
   // Get the default level of parallelism to use in the cluster, as a hint for sizing jobs.
   def defaultParallelism(): Int
 
+  /**
+   * Identity of the executor and retry partition reserved for OOM recovery.
+   * Allocation derives covered tasks from its own listener state, including during launch lag.
+   * Dynamic allocation reads this while holding its own lock, so it must not take scheduler locks.
+   */
+  def oomRetryReservationInfo: Option[OomRetryReservationInfo] = None
+
   /**
    * Update metrics for in-progress tasks and executor metrics, and let the master know that the
    * BlockManager is still alive. Return true if the driver knows about the given block manager.