diff --git a/core/src/main/scala/org/apache/spark/ExecutorAllocationManager.scala b/core/src/main/scala/org/apache/spark/ExecutorAllocationManager.scala index 9ec5da71406f7..2d857f01b8c67 100644 --- a/core/src/main/scala/org/apache/spark/ExecutorAllocationManager.scala +++ b/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,6 +974,16 @@ 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 @@ -950,7 +991,8 @@ private[spark] class ExecutorAllocationManager( // (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 diff --git a/core/src/main/scala/org/apache/spark/SparkContext.scala b/core/src/main/scala/org/apache/spark/SparkContext.scala index 5c8c30e2807c7..743b15050d5b1 100644 --- a/core/src/main/scala/org/apache/spark/SparkContext.scala +++ b/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 } diff --git a/core/src/main/scala/org/apache/spark/TaskEndReason.scala b/core/src/main/scala/org/apache/spark/TaskEndReason.scala index 91ed190d7e2ff..9ceff665a2ad4 100644 --- a/core/src/main/scala/org/apache/spark/TaskEndReason.scala +++ b/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" diff --git a/core/src/main/scala/org/apache/spark/internal/config/package.scala b/core/src/main/scala/org/apache/spark/internal/config/package.scala index 87373d777d937..a87f49bf12a71 100644 --- a/core/src/main/scala/org/apache/spark/internal/config/package.scala +++ b/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") diff --git a/core/src/main/scala/org/apache/spark/scheduler/ExecutorLossReason.scala b/core/src/main/scala/org/apache/spark/scheduler/ExecutorLossReason.scala index fb6a62551fa44..9640c1a140c3f 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/ExecutorLossReason.scala +++ b/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 = { diff --git a/core/src/main/scala/org/apache/spark/scheduler/ExecutorResourcesAmounts.scala b/core/src/main/scala/org/apache/spark/scheduler/ExecutorResourcesAmounts.scala index caad1cb6780e9..d8b2cfe0bebd6 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/ExecutorResourcesAmounts.scala +++ b/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 diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskScheduler.scala index 1e6de9ef46f34..ff981318a0eee 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskScheduler.scala +++ b/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. diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala index 4600993d2664e..e22fb25d41cda 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala @@ -19,7 +19,7 @@ package org.apache.spark.scheduler import java.nio.ByteBuffer import java.util.{Properties, TimerTask} -import java.util.concurrent.{ConcurrentHashMap, TimeUnit} +import java.util.concurrent.{ConcurrentHashMap, ScheduledFuture, TimeUnit} import java.util.concurrent.atomic.AtomicLong import scala.collection.mutable @@ -121,6 +121,26 @@ private[spark] class TaskSchedulerImpl( // the exact BigDecimal the user configured. val CPUS_PER_TASK = conf.get(config.CPUS_PER_TASK) + private[scheduler] val oomRetryEnabled = conf.get(SCHEDULER_OOM_RETRY_ENABLED) + private[scheduler] val oomRetryIsolationTimeoutMs = + conf.get(SCHEDULER_OOM_RETRY_ISOLATION_TIMEOUT) + + // One application-wide reservation, held across offer rounds while an executor drains and + // until the isolated attempt exits. CPU accounting and TaskContext.cpus remain unchanged. + private case class OomRetryReservation( + taskSet: TaskSetManager, + index: Int, + executorId: String, + taskId: Option[Long] = None) + + private var oomRetryReservation: Option[OomRetryReservation] = None + private var oomRetryWakeup: Option[ScheduledFuture[_]] = None + + @volatile private var oomRetryReservationSnapshot: Option[OomRetryReservationInfo] = None + + override def oomRetryReservationInfo: Option[OomRetryReservationInfo] = + oomRetryReservationSnapshot + // TaskSetManagers are not thread safe, so any access to one should be synchronized // on this class. Protected by `this` private val taskSetsByStageIdAndAttempt = new HashMap[Int, HashMap[Int, TaskSetManager]] @@ -400,6 +420,9 @@ private[spark] class TaskSchedulerImpl( log"${MDC(LogKeys.STAGE_ATTEMPT_ID, tsm.taskSet.stageAttemptId)} was cancelled") } } + if (refreshOomRetryReservation()) { + backend.reviveOffers() + } } override def killTaskAttempt( @@ -428,6 +451,8 @@ private[spark] class TaskSchedulerImpl( * cleaned up. */ def taskSetFinished(manager: TaskSetManager): Unit = synchronized { + val releasedOomReservation = oomRetryReservation.exists(_.taskSet == manager) + setOomRetryReservation(oomRetryReservation.filterNot(_.taskSet == manager)) taskSetsByStageIdAndAttempt.get(manager.taskSet.stageId).foreach { taskSetsForStage => taskSetsForStage -= manager.taskSet.stageAttemptId if (taskSetsForStage.isEmpty) { @@ -439,6 +464,185 @@ private[spark] class TaskSchedulerImpl( logInfo(log"Removed TaskSet " + manager.taskSet.logId + log" whose tasks have all completed, from pool ${MDC(LogKeys.POOL_NAME, manager.parent.name)}" ) + if (releasedOomReservation) { + // A cancelled pending retry has no task exit to wake work held behind its reservation. + backend.reviveOffers() + } + } + + private def setOomRetryReservation(reservation: Option[OomRetryReservation]): Unit = { + oomRetryReservation = reservation + // Publish an immutable snapshot without calling into the allocation manager: it can already + // hold its lock while acquiring the scheduler lock through the backend. + oomRetryReservationSnapshot = reservation.map { r => + val taskSet = r.taskSet.taskSet + OomRetryReservationInfo(taskSet.resourceProfileId, r.executorId, + taskSet.stageId, taskSet.stageAttemptId, r.index) + } + } + + private def refreshOomRetryReservation(): Boolean = { + val hadReservation = oomRetryReservation.nonEmpty + setOomRetryReservation(oomRetryReservation.filter { reservation => + reservation.taskId match { + // Cancellation or another attempt's success can precede the task's exit. Keep its + // executor isolated until the terminal status update releases the running task. + case Some(tid) => taskIdToExecutorId.get(tid).contains(reservation.executorId) + case None => + isExecutorAlive(reservation.executorId) && + executorIdToHost.get(reservation.executorId).exists { host => + !healthTrackerOpt.exists { tracker => + tracker.isExecutorExcluded(reservation.executorId) || tracker.isNodeExcluded(host) + } && reservation.taskSet.canRunOomRetry( + reservation.index, reservation.executorId, host) + } && reservation.taskSet.oomRetryNeedsIsolation(reservation.index) + } + }) + hadReservation && oomRetryReservation.isEmpty + } + + /** + * Give OOM retries idle capacity before ordinary placement. Repeated OOMs reserve at most one + * executor across offer rounds, allowing existing tasks to drain without backfilling it. Only + * this recovery pass ignores locality preferences; resource requirements and exclusions remain + * mandatory. A bounded wait in TaskSetManager makes an unavailable reservation fall back to + * ordinary scheduling instead of turning an OOM into an indefinitely pending retry. + */ + private def scheduleOomRetries( + taskSets: ArrayBuffer[TaskSetManager], + offers: IndexedSeq[WorkerOffer], + availableCpus: Array[BigDecimal], + availableResources: Array[ExecutorResourcesAmounts], + tasks: IndexedSeq[ArrayBuffer[TaskDescription]]): Set[TaskSetManager] = { + refreshOomRetryReservation() + val launchedTaskSets = new HashSet[TaskSetManager] + def taskCpusFor(taskSet: TaskSetManager): BigDecimal = { + val profile = sc.resourceProfileManager + .resourceProfileFromId(taskSet.taskSet.resourceProfileId) + ResourceProfile.getTaskCpusOrDefaultForProfile(profile, conf) + } + + // A compatible profile does not guarantee that every executor has enough total resources. + // Once a reserved executor is idle, its offer tells us whether the retry can actually fit. + setOomRetryReservation(oomRetryReservation.filterNot { reservation => + reservation.taskId.isEmpty && offers.indices.exists { i => + offers(i).executorId == reservation.executorId && + !isExecutorBusy(reservation.executorId) && + resourcesMeetTaskRequirements( + reservation.taskSet, taskCpusFor(reservation.taskSet), + availableCpus(i), availableResources(i)).isEmpty + } + }) + + def eligible(taskSet: TaskSetManager, index: Int, i: Int): Boolean = { + val offer = offers(i) + isExecutorAlive(offer.executorId) && + sc.resourceProfileManager.canBeScheduled( + taskSet.taskSet.resourceProfileId, offer.resourceProfileId) && + taskSet.canRunOomRetry(index, offer.executorId, offer.host) + } + + def launch(taskSet: TaskSetManager, index: Int, i: Int): Option[Long] = { + val offer = offers(i) + if (!eligible(taskSet, index, i) || isExecutorBusy(offer.executorId)) { + return None + } + resourcesMeetTaskRequirements( + taskSet, taskCpusFor(taskSet), availableCpus(i), availableResources(i)).flatMap { + assignments => + val taskCpus = taskCpusFor(taskSet) + try { + taskSet.resourceOfferOomRetry( + index, offer.executorId, offer.host, taskCpus, assignments) + .map { task => + tasks(i) += task + addRunningTask(task.taskId, offer.executorId, taskSet) + launchedTaskSets += taskSet + availableCpus(i) -= task.cpus + availableResources(i).acquire(task.resources) + task.taskId + } + } catch { + case e: TaskNotSerializableException => + // prepareLaunchingTask already aborts the task set on serialization failure. + logError("Failed to serialize an OOM retry", e) + None + } + } + } + + if (oomRetryReservation.isEmpty) { + for { + taskSet <- taskSets + index <- taskSet.pendingOomRetries + if oomRetryReservation.isEmpty && taskSet.oomRetryNeedsIsolation(index) + } { + val candidates = offers.indices.filter { i => + eligible(taskSet, index, i) && + (isExecutorBusy(offers(i).executorId) || + resourcesMeetTaskRequirements( + taskSet, taskCpusFor(taskSet), availableCpus(i), availableResources(i)) + .isDefined) + } + // A busy executor may currently offer no CPUs or custom resources. Check its profile + // now and its actual free resources again after it drains, before launching anything. + candidates.sortBy(i => executorIdToRunningTaskIds(offers(i).executorId).size) + .headOption.foreach { i => + setOomRetryReservation(Some( + OomRetryReservation(taskSet, index, offers(i).executorId))) + logInfo(log"Reserving executor ${MDC(LogKeys.EXECUTOR_ID, offers(i).executorId)} for " + + log"OOM retry of task ${MDC(TASK_INDEX, index)} in stage " + taskSet.taskSet.logId) + } + } + } + + oomRetryReservation.filter(_.taskId.isEmpty).foreach { reservation => + offers.indices.find(i => offers(i).executorId == reservation.executorId).foreach { i => + launch(reservation.taskSet, reservation.index, i).foreach { tid => + setOomRetryReservation(Some(reservation.copy(taskId = Some(tid)))) + } + } + } + + for { + taskSet <- taskSets + index <- taskSet.pendingOomRetries + if !taskSet.oomRetryNeedsIsolation(index) + } { + offers.indices.iterator.filter { i => + !isExecutorBusy(offers(i).executorId) && + !oomRetryReservation.exists(_.executorId == offers(i).executorId) + }.exists(i => launch(taskSet, index, i).isDefined) + } + + // Capture deadlines before refreshing and masking: if one expires during this pass, either + // its reservation is released below or the captured delay still guarantees another offer. + val nextDelay = taskSets.iterator.flatMap { taskSet => + taskSet.pendingOomRetries.iterator.map(taskSet.oomRetryIsolationTimeRemaining) + }.filter(_ > 0).reduceOption(_ min _) + + // A failed launch may have aborted its task set. Missing partial offers must not release an + // otherwise valid reservation. Mask it for every task set, including barrier slot counting. + refreshOomRetryReservation() + oomRetryReservation.foreach { reservation => + offers.indices.filter(i => offers(i).executorId == reservation.executorId) + .foreach { i => + availableCpus(i) = 0 + availableResources(i) = ExecutorResourcesAmounts.empty + } + } + + // Local backends do not periodically revive offers. Keep one wakeup for the earliest + // pending deadline, including retries waiting behind another task's reservation. + oomRetryWakeup.foreach(_.cancel(false)) + oomRetryWakeup = nextDelay.map { delay => + starvationTimer.schedule(new Runnable { + override def run(): Unit = Utils.tryLogNonFatalError { + backend.reviveOffers() + } + }, delay, TimeUnit.MILLISECONDS) + } + launchedTaskSets.toSet } /** @@ -619,6 +823,12 @@ private[spark] class TaskSchedulerImpl( val availableCpus = shuffledOffers.map(o => o.cores).toArray val resourceProfileIds = shuffledOffers.map(o => o.resourceProfileId).toArray val sortedTaskSets = rootPool.getSortedTaskSetQueue + val oomRetryTaskSets = if (oomRetryEnabled) { + scheduleOomRetries( + sortedTaskSets, shuffledOffers, availableCpus, availableResources, tasks) + } else { + Set.empty[TaskSetManager] + } for (taskSet <- sortedTaskSets) { logDebug("parentName: %s, name: %s, runningTasks: %s".format( taskSet.parent.name, taskSet.name, taskSet.runningTasks)) @@ -635,8 +845,16 @@ private[spark] class TaskSchedulerImpl( // value is -1 val numBarrierSlotsAvailable = if (taskSet.isBarrier) { val rpId = taskSet.taskSet.resourceProfileId - val resAmounts = availableResources.map(_.resourceAddressAmount) - calculateAvailableSlots(this, conf, rpId, resourceProfileIds, availableCpus, resAmounts) + val profile = sc.resourceProfileManager.resourceProfileFromId(rpId) + val taskCpus = ResourceProfile.getTaskCpusOrDefaultForProfile(profile, conf) + shuffledOffers.indices.iterator + .filter(i => sc.resourceProfileManager.canBeScheduled(rpId, resourceProfileIds(i))) + .map { i => + // Recovery may already have consumed whole or fractional custom resources. + val maxTasks = math.min( + ResourceProfile.numTasksBasedOnCores(availableCpus(i), taskCpus), taskSet.numTasks) + availableResources(i).availableTaskSlots(profile, maxTasks) + }.foldLeft(0L)(_ + _).min(taskSet.numTasks.toLong).toInt } else { -1 } @@ -650,7 +868,8 @@ private[spark] class TaskSchedulerImpl( log"${MDC(LogKeys.TASK_SET_NAME, taskSet.numTasks)} slots, while the total " + log"number of available slots is ${MDC(LogKeys.NUM_SLOTS, numBarrierSlotsAvailable)}.") } else { - var launchedAnyTask = false + // Recovery launches also clear unschedulable-task-set timers, without advancing locality. + var launchedAnyTask = oomRetryTaskSets.contains(taskSet) var noDelaySchedulingRejects = true var globalMinLocality: Option[TaskLocality] = None for (currentMaxLocality <- taskSet.myLocalityLevels) { @@ -742,6 +961,12 @@ private[spark] class TaskSchedulerImpl( if (launchedAnyTask && taskSet.isBarrier) { val barrierPendingLaunchTasks = taskSet.barrierPendingLaunchTasks.values.toArray + def releaseAssignedResources(): Unit = { + barrierPendingLaunchTasks.foreach { task => + availableCpus(task.assignedOfferIndex) += task.assignedCores + availableResources(task.assignedOfferIndex).release(task.assignedResources) + } + } // Check whether the barrier tasks are partially launched. if (barrierPendingLaunchTasks.length != taskSet.numTasks) { if (legacyLocalityWaitReset) { @@ -761,7 +986,6 @@ private[spark] class TaskSchedulerImpl( log"to get rid of this error." logWarning(logMsg) taskSet.abort(logMsg.message) - throw SparkCoreErrors.sparkError(logMsg.message) } else { val curTime = clock.getTimeMillis() if (curTime - taskSet.lastResourceOfferFailLogTime > @@ -771,47 +995,64 @@ private[spark] class TaskSchedulerImpl( taskSet.lastResourceOfferFailLogTime = curTime } barrierPendingLaunchTasks.foreach { task => - // revert all assigned resources - availableCpus(task.assignedOfferIndex) = - availableCpus(task.assignedOfferIndex) + task.assignedCores - availableResources(task.assignedOfferIndex).release( - task.assignedResources) // re-add the task to the schedule pending list taskSet.addPendingTask(task.index) } } + releaseAssignedResources() } else { // All tasks are able to launch in this barrier task set. Let's do // some preparation work before launching them. val launchTime = clock.getTimeMillis() - val addressesWithDescs = barrierPendingLaunchTasks.map { task => - val taskDesc = taskSet.prepareLaunchingTask( - task.execId, - task.host, - task.index, - task.taskLocality, - false, - task.assignedCores, - task.assignedResources, - launchTime) - addRunningTask(taskDesc.taskId, taskDesc.executorId, taskSet) - tasks(task.assignedOfferIndex) += taskDesc - shuffledOffers(task.assignedOfferIndex).address.get -> taskDesc - } - - // materialize the barrier coordinator. - maybeInitBarrierCoordinator() - - // Update the taskInfos into all the barrier task properties. - val addressesStr = addressesWithDescs - // Addresses ordered by partitionId - .sortBy(_._2.partitionId) - .map(_._1) - .mkString(",") - addressesWithDescs.foreach(_._2.properties.setProperty("addresses", addressesStr)) + val addressesWithDescs = new ArrayBuffer[(String, TaskDescription)] + try { + barrierPendingLaunchTasks.foreach { task => + val taskDesc = taskSet.prepareLaunchingTask( + task.execId, + task.host, + task.index, + task.taskLocality, + false, + task.assignedCores, + task.assignedResources, + launchTime) + addressesWithDescs += + shuffledOffers(task.assignedOfferIndex).address.get -> taskDesc + } - logInfo(log"Successfully scheduled all the ${MDC(LogKeys.NUM_TASKS, addressesWithDescs.length)} " + - log"tasks for barrier stage ${MDC(LogKeys.STAGE_ID, taskSet.stageId)}.") + // materialize the barrier coordinator. + maybeInitBarrierCoordinator() + + // Update the taskInfos into all the barrier task properties. + val addressesStr = addressesWithDescs + // Addresses ordered by partitionId + .sortBy(_._2.partitionId) + .map(_._1) + .mkString(",") + addressesWithDescs.foreach(_._2.properties.setProperty("addresses", addressesStr)) + + // Publish the group only after every task is prepared. A serialization failure + // must not launch a partial barrier group or discard earlier ordinary/OOM tasks. + barrierPendingLaunchTasks.zip(addressesWithDescs).foreach { + case (task, (_, taskDesc)) => + addRunningTask(taskDesc.taskId, taskDesc.executorId, taskSet) + tasks(task.assignedOfferIndex) += taskDesc + } + logInfo(log"Successfully scheduled all the " + + log"${MDC(LogKeys.NUM_TASKS, addressesWithDescs.length)} " + + log"tasks for barrier stage ${MDC(LogKeys.STAGE_ID, taskSet.stageId)}.") + } catch { + case e: TaskNotSerializableException => + // prepareLaunchingTask already aborted this task set. Balance the task-start + // bookkeeping of earlier preparations without sending them to an executor. + releaseAssignedResources() + addressesWithDescs.foreach { case (_, taskDesc) => + taskSet.handleFailedTask(taskDesc.taskId, TaskState.KILLED, + TaskKilled("Barrier task serialization failed")) + } + logError(log"Failed to serialize barrier task set " + + log"${MDC(TASK_SET_NAME, taskSet.name)}", e) + } } taskSet.barrierPendingLaunchTasks.clear() } @@ -990,6 +1231,9 @@ private[spark] class TaskSchedulerImpl( taskSetsByStageIdAndAttempt.get(stageId).foreach(_.values.filter(!_.isZombie).foreach { tsm => tsm.markPartitionCompleted(partitionId) }) + if (refreshOomRetryReservation()) { + backend.reviveOffers() + } } def error(message: String): Unit = { @@ -1034,6 +1278,11 @@ private[spark] class TaskSchedulerImpl( barrierCoordinator.stop() } } + synchronized { + oomRetryWakeup.foreach(_.cancel(false)) + oomRetryWakeup = None + setOomRetryReservation(None) + } ThreadUtils.shutdown(starvationTimer) ThreadUtils.shutdown(abortTimer) } @@ -1144,6 +1393,7 @@ private[spark] class TaskSchedulerImpl( taskIdToExecutorId.remove(tid).foreach { executorId => executorIdToRunningTaskIds.get(executorId).foreach { _.remove(tid) } } + setOomRetryReservation(oomRetryReservation.filterNot(_.taskId.contains(tid))) } /** @@ -1152,6 +1402,7 @@ private[spark] class TaskSchedulerImpl( * of any running tasks, since the loss reason defines whether we'll fail those tasks. */ private def removeExecutor(executorId: String, reason: ExecutorLossReason): Unit = { + setOomRetryReservation(oomRetryReservation.filterNot(_.executorId == executorId)) // The tasks on the lost executor may not send any more status updates (because the executor // has been lost), so they should be cleaned up here. executorIdToRunningTaskIds.remove(executorId).foreach { taskIds => diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala index 9920d7ad971fb..9242e452792f5 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala @@ -135,6 +135,29 @@ private[spark] class TaskSetManager( val successful = new Array[Boolean](numTasks) private val numFailures = new Array[Int](numTasks) + // Sparse state for OOM-affected partitions: failure count and time of the latest OOM. The + // deadline also bounds retries waiting behind another task's executor reservation. + private val oomRetries = new HashMap[Int, (Int, Long)] + + private[scheduler] def isPendingOomRetry(index: Int): Boolean = { + !isZombie && oomRetries.contains(index) && !successful(index) && copiesRunning(index) == 0 + } + + private[scheduler] def pendingOomRetries: Seq[Int] = { + oomRetries.keysIterator.filter(isPendingOomRetry).toSeq.sorted + } + + private[scheduler] def oomRetryNeedsIsolation(index: Int): Boolean = { + oomRetryIsolationTimeRemaining(index) > 0 + } + + private[scheduler] def oomRetryIsolationTimeRemaining(index: Int): Long = { + if (!sched.oomRetryEnabled) return 0L + oomRetries.get(index).collect { case (failures, failedAt) if failures >= 2 => + sched.oomRetryIsolationTimeoutMs - (clock.getTimeMillis() - failedAt) + }.getOrElse(0L) + } + // Add the tid of task into this HashSet when the task is killed by other attempt tasks. // This happened while we set the `spark.speculation` to true. The task killed by others // should not resubmit while executor lost. @@ -345,7 +368,8 @@ private[spark] class TaskSetManager( indexOffset -= 1 val index = list(indexOffset) if (!isTaskExcludededOnExecOrNode(index, execId, host) && - !(speculative && hasAttemptOnHost(index, host))) { + !(speculative && (hasAttemptOnHost(index, host) || oomRetries.contains(index))) && + !oomRetryNeedsIsolation(index)) { // This should almost always be list.trimEnd(1) to remove tail list.remove(indexOffset) // Speculatable task should only be launched when at most one copy of the @@ -374,6 +398,34 @@ private[spark] class TaskSetManager( } } + private def isOfferExcluded(execId: String, host: String): Boolean = { + taskSetExcludelistHelperOpt.exists { excludeList => + excludeList.isNodeExcludedForTaskSet(host) || + excludeList.isExecutorExcludedForTaskSet(execId) + } + } + + private[scheduler] def canRunOomRetry(index: Int, execId: String, host: String): Boolean = { + isPendingOomRetry(index) && !isOfferExcluded(execId, host) && + !isTaskExcludededOnExecOrNode(index, execId, host) + } + + /** OOM recovery favors memory headroom over preferred locations, without changing task CPUs. */ + private[scheduler] def resourceOfferOomRetry( + index: Int, + execId: String, + host: String, + taskCpus: BigDecimal, + taskResources: Map[String, Map[String, Long]]): Option[TaskDescription] = { + if (canRunOomRetry(index, execId, host)) { + // Other pending-task lists are cleaned lazily, just as for the ordinary dequeue path. + Some(prepareLaunchingTask(execId, host, index, TaskLocality.ANY, false, + taskCpus, taskResources, clock.getTimeMillis())) + } else { + None + } + } + /** * Dequeue a pending task for a given node and return its index and locality level. * Only search for tasks matching the given locality constraint. @@ -476,11 +528,7 @@ private[spark] class TaskSetManager( taskResourceAssignments: Map[String, Map[String, Long]] = Map.empty) : (Option[TaskDescription], Boolean, Int) = { - val offerExcluded = taskSetExcludelistHelperOpt.exists { excludeList => - excludeList.isNodeExcludedForTaskSet(host) || - excludeList.isExecutorExcludedForTaskSet(execId) - } - if (!isZombie && !offerExcluded) { + if (!isZombie && !isOfferExcluded(execId, host)) { val curTime = clock.getTimeMillis() var allowedLocality = maxLocality @@ -885,6 +933,7 @@ private[spark] class TaskSetManager( // Mark successful and stop if all the tasks have succeeded. successful(index) = true numFailures(index) = 0 + oomRetries.remove(index) if (tasksSuccessful == numTasks) { isZombie = true } @@ -977,6 +1026,7 @@ private[spark] class TaskSetManager( tasksSuccessful += 1 successful(index) = true numFailures(index) = 0 + oomRetries.remove(index) if (tasksSuccessful == numTasks) { isZombie = true } @@ -1005,6 +1055,17 @@ private[spark] class TaskSetManager( info.markFinished(state, clock.getTimeMillis()) val index = info.index copiesRunning(index) -= 1 + val isOom = reason match { + case e: ExceptionFailure => e.isOutOfMemoryError + case e: ExecutorLostFailure => e.exitCausedByApp && e.isOutOfMemoryError + case _ => false + } + if (sched.oomRetryEnabled && !isBarrier && !taskSet.isPipelined && + !successful(index) && isOom) { + val failures = oomRetries.get(index).map(_._1).getOrElse(0) + 1 + oomRetries(index) = (failures, clock.getTimeMillis()) + speculatableTasks -= index + } var accumUpdates: Seq[AccumulatorV2[_, _]] = Seq.empty var metricPeaks: Array[Long] = Array.empty val failureReason = log"Lost ${MDC(TASK_NAME, taskName(tid))} " + @@ -1291,8 +1352,12 @@ private[spark] class TaskSetManager( // 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))) + val failure = ExecutorLostFailure(info.executorId, exitCausedByApp, Some(reason.toString)) + failure.isOutOfMemoryError = exitCausedByApp && (reason match { + case e: ExecutorExited => e.isOutOfMemoryError + case _ => false + }) + handleFailedTask(tid, TaskState.FAILED, failure) } } // recalculate valid locality levels and waits when executor is lost @@ -1311,7 +1376,8 @@ private[spark] class TaskSetManager( for (tid <- runningTasksSet) { val info = taskInfos(tid) val index = info.index - if (!successful(index) && copiesRunning(index) == 1 && !speculatableTasks.contains(index)) { + if (!successful(index) && copiesRunning(index) == 1 && !speculatableTasks.contains(index) && + !oomRetries.contains(index)) { val runtimeMs = info.timeRunning(currentTimeMillis) def checkMaySpeculate(): Boolean = { diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala index d56da893f86fc..25f0c8d47e1e0 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala @@ -26,7 +26,7 @@ import scala.concurrent.Future import com.google.common.cache.CacheBuilder -import org.apache.spark.{ExecutorAllocationClient, SparkEnv, TaskState, VersionedCredentials} +import org.apache.spark.{ExecutorAllocationClient, SparkEnv, TaskKilled, TaskState, VersionedCredentials} import org.apache.spark.deploy.SparkHadoopUtil import org.apache.spark.deploy.security.UserCredentialManager import org.apache.spark.errors.SparkCoreErrors @@ -468,7 +468,15 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp s"${RPC_MESSAGE_MAX_SIZE.key} (%d bytes). Consider increasing " + s"${RPC_MESSAGE_MAX_SIZE.key} or using broadcast variables for large values." msg = msg.format(task.taskId, task.index, serializedTask.limit(), maxRpcMessageSize) - taskSetMgr.abort(msg) + try { + taskSetMgr.abort(msg) + } finally { + // The task was registered with the scheduler but will never report its status. + // Do not send a backend StatusUpdate, since no resources were allocated above. + val reason = scheduler.sc.env.closureSerializer.newInstance() + .serialize(TaskKilled(msg)) + scheduler.statusUpdate(task.taskId, TaskState.KILLED, reason) + } } catch { case e: Exception => logError("Exception in error callback", e) } diff --git a/core/src/main/scala/org/apache/spark/util/JsonProtocol.scala b/core/src/main/scala/org/apache/spark/util/JsonProtocol.scala index 2d5a978b49e9b..23674b9643e05 100644 --- a/core/src/main/scala/org/apache/spark/util/JsonProtocol.scala +++ b/core/src/main/scala/org/apache/spark/util/JsonProtocol.scala @@ -711,14 +711,20 @@ private[spark] object JsonProtocol extends JsonUtils { g.writeStringField("Full Stack Trace", exceptionFailure.fullStackTrace) g.writeFieldName("Accumulator Updates") accumulablesToJson(exceptionFailure.accumUpdates, g) + if (exceptionFailure.isOutOfMemoryError) { + g.writeBooleanField("Out Of Memory Error", true) + } case taskCommitDenied: TaskCommitDenied => g.writeNumberField("Job ID", taskCommitDenied.jobID) g.writeNumberField("Partition ID", taskCommitDenied.partitionID) g.writeNumberField("Attempt Number", taskCommitDenied.attemptNumber) - case ExecutorLostFailure(executorId, exitCausedByApp, reason) => + case executorLostFailure @ ExecutorLostFailure(executorId, exitCausedByApp, reason) => g.writeStringField("Executor ID", executorId) g.writeBooleanField("Exit Caused By App", exitCausedByApp) reason.foreach(g.writeStringField("Loss Reason", _)) + if (executorLostFailure.isOutOfMemoryError) { + g.writeBooleanField("Out Of Memory Error", true) + } case taskKilled: TaskKilled => g.writeStringField("Kill Reason", taskKilled.reason) g.writeFieldName("Accumulator Updates") @@ -1457,7 +1463,11 @@ private[spark] object JsonProtocol extends JsonUtils { .getOrElse(taskMetricsFromJson(json.get("Metrics")).accumulators().map(acc => { acc.toInfoUpdate }).toArray.toImmutableArraySeq) - ExceptionFailure(className, description, stackTrace, fullStackTrace, None, accumUpdates) + val failure = + ExceptionFailure(className, description, stackTrace, fullStackTrace, None, accumUpdates) + failure.isOutOfMemoryError = jsonOption(json.get("Out Of Memory Error")) + .map(_.extractBoolean).getOrElse(failure.isOutOfMemoryError) + failure case `taskResultLost` => TaskResultLost case `taskKilled` => // The "Kill Reason" field was added in Spark 2.2.0: @@ -1480,10 +1490,13 @@ private[spark] object JsonProtocol extends JsonUtils { val exitCausedByApp = jsonOption(json.get("Exit Caused By App")).map(_.extractBoolean) val executorId = jsonOption(json.get("Executor ID")).map(_.asText) val reason = jsonOption(json.get("Loss Reason")).map(_.asText) - ExecutorLostFailure( + val failure = ExecutorLostFailure( executorId.getOrElse("Unknown"), exitCausedByApp.getOrElse(true), reason) + failure.isOutOfMemoryError = jsonOption(json.get("Out Of Memory Error")) + .exists(_.extractBoolean) + failure case `executorShutdownFailure` => val executorId = json.get("Executor ID").extractString ExecutorShutdownFailure(executorId) diff --git a/core/src/test/scala/org/apache/spark/ExecutorAllocationManagerSuite.scala b/core/src/test/scala/org/apache/spark/ExecutorAllocationManagerSuite.scala index 32289e74ccacd..1444afbaddda4 100644 --- a/core/src/test/scala/org/apache/spark/ExecutorAllocationManagerSuite.scala +++ b/core/src/test/scala/org/apache/spark/ExecutorAllocationManagerSuite.scala @@ -513,6 +513,419 @@ class ExecutorAllocationManagerSuite extends SparkFunSuite { assert(numExecutorsToAddForDefaultProfile(manager) === 1) } + for { + enabled <- Seq(false, true) + (failureName, reason, attributedOom) <- Seq( + ("OOM exception", new ExceptionFailure( + new RuntimeException("wrapped", new OutOfMemoryError("test OOM")), Nil), true), + ("ordinary exception", new ExceptionFailure( + new RuntimeException("OutOfMemoryError in an ordinary exception's message"), Nil), false), + ("OOM executor loss", { + val failure = ExecutorLostFailure("a", exitCausedByApp = true, Some("OOMKilled")) + failure.isOutOfMemoryError = true + failure + }, true), + ("unattributed OOM executor loss", { + val failure = ExecutorLostFailure("a", exitCausedByApp = false, Some("OOMKilled")) + failure.isOutOfMemoryError = true + failure + }, false), + ("executor loss without typed OOM", ExecutorLostFailure( + "a", exitCausedByApp = true, Some("OOMKilled")), false)) + } { + test(s"OOM recovery speculative demand after $failureName (enabled=$enabled)") { + val clock = new ManualClock(1) + val conf = createConf(0, 5, 0) + .set(config.EXECUTOR_CORES, 4) + .set(config.SCHEDULER_OOM_RETRY_ENABLED, enabled) + val manager = createManager(conf, clock = clock) + val rpId = DEFAULT_RESOURCE_PROFILE_ID + val revokeSpeculation = enabled && attributedOom + post(SparkListenerStageSubmitted(createStageInfo(0, 1))) + onExecutorAdded(manager, "a", rpManager.defaultResourceProfile) + clock.advance(schedulerBacklogTimeout * 1000) + schedule(manager) + assert(numExecutorsTargetForDefaultProfileId(manager) === 1) + + val original = createTaskInfo(0, 0, "a") + post(SparkListenerTaskStart(0, 0, original)) + assert(addTime(manager) === NOT_SET) + post(speculativeTaskSubmitEventFromTaskIndex(0, taskIndex = 0)) + clock.advance(schedulerBacklogTimeout * 1000) + schedule(manager) + assert(numExecutorsTargetForDefaultProfileId(manager) === 2) + + original.markFinished(TaskState.FAILED, clock.getTimeMillis()) + post(SparkListenerTaskEnd(0, 0, null, reason, original, new ExecutorMetrics, null)) + assert(manager.listener.pendingSpeculativeTasksPerResourceProfile(rpId) === + (if (revokeSpeculation) 0 else 1)) + // Removing revoked speculation must leave the ordinary failed task backlogged. + assert(manager.listener.pendingTasksPerResourceProfile(rpId) === 1) + assert(addTime(manager) !== NOT_SET) + + val retryExecutor = if (reason.isInstanceOf[ExecutorLostFailure]) { + onExecutorRemoved(manager, "a") + onExecutorAdded(manager, "replacement", rpManager.defaultResourceProfile) + "replacement" + } else { + "a" + } + val retry = createTaskInfo(1, 0, retryExecutor) + post(SparkListenerTaskStart(0, 0, retry)) + onExecutorAdded(manager, "b", rpManager.defaultResourceProfile) + assert(manager.listener.pendingTasksPerResourceProfile(rpId) === 0) + assert(totalRunningTasksPerResourceProfile(manager) === 1) + assert((addTime(manager) == NOT_SET) === revokeSpeculation) + + // With the retry running, only valid speculative demand may retain the idle executor. + clock.advance(executorIdleTimeout * 1000 + 1) + assert(manager.executorMonitor.timedOutExecutors().map(_._1) === Seq("b")) + schedule(manager) + assert(numExecutorsTargetForDefaultProfileId(manager) === + (if (revokeSpeculation) 1 else 2)) + assert(executorsPendingToRemove(manager) === + (if (revokeSpeculation) Set("b") else Set.empty[String])) + + retry.markFinished(TaskState.FINISHED, clock.getTimeMillis()) + post(SparkListenerTaskEnd(0, 0, null, Success, retry, new ExecutorMetrics, null)) + assert(manager.listener.pendingSpeculativeTasksPerResourceProfile(rpId) === 0) + schedule(manager) + assert(numExecutorsTargetForDefaultProfileId(manager) === 0) + assert(executorsPendingToRemove(manager).contains("b")) + } + } + + Seq(1.0, 0.1).foreach { allocationRatio => + test(s"OOM retry isolation preserves ordinary executor demand (ratio=$allocationRatio)") { + val clock = new ManualClock(1) + val conf = createConf(0, 5, 1) + .set(config.EXECUTOR_CORES, 4) + .set(config.SCHEDULER_OOM_RETRY_ENABLED, true) + .set(config.DYN_ALLOCATION_EXECUTOR_ALLOCATION_RATIO, allocationRatio) + var reservation: Option[OomRetryReservationInfo] = None + val manager = createManager(conf, clock = clock, oomRetryReservationInfo = () => reservation) + val rpId = DEFAULT_RESOURCE_PROFILE_ID + onExecutorAdded(manager, "reserved", rpManager.defaultResourceProfile) + post(SparkListenerStageSubmitted(createStageInfo(0, 1))) + post(SparkListenerTaskStart(0, 0, createTaskInfo(0, 0, "reserved"))) + post(SparkListenerStageSubmitted(createStageInfo(1, 3))) + + assert(manager.listener.pendingTasksPerResourceProfile(rpId) === 3) + assert(totalRunningTasksPerResourceProfile(manager) === 1) + assert(manager.listener.pendingUnschedulableTaskSetsPerResourceProfile(rpId) === 0) + assert(manager.maxNumExecutorsNeededPerResourceProfile(rpId) === 1) + + // The isolated retry cannot share its executor's other three slots with ordinary tasks. + reservation = Some(OomRetryReservationInfo(rpId, "reserved", 0, 0, 0)) + assert(manager.maxNumExecutorsNeededPerResourceProfile(rpId) === 2) + clock.advance(schedulerBacklogTimeout * 1000) + schedule(manager) + assert(numExecutorsTargetForDefaultProfileId(manager) === 2) + + // Releasing the reservation restores the normal target without changing task demand. + reservation = None + schedule(manager) + assert(manager.maxNumExecutorsNeededPerResourceProfile(rpId) === 1) + assert(numExecutorsTargetForDefaultProfileId(manager) === 1) + } + } + + test("OOM retry isolation accounts for tasks already draining on the reserved executor") { + val clock = new ManualClock(1) + val conf = createConf(0, 5, 1) + .set(config.EXECUTOR_CORES, 4) + .set(config.SCHEDULER_OOM_RETRY_ENABLED, true) + val rpId = DEFAULT_RESOURCE_PROFILE_ID + val reservation = Some(OomRetryReservationInfo(rpId, "reserved", 0, 0, 4)) + val manager = createManager(conf, clock = clock, oomRetryReservationInfo = () => reservation) + onExecutorAdded(manager, "reserved", rpManager.defaultResourceProfile) + post(SparkListenerStageSubmitted(createStageInfo(0, 5))) + val draining = (0 until 4).map { index => createTaskInfo(index, index, "reserved") } + draining.foreach { info => post(SparkListenerTaskStart(0, 0, info)) } + + // Four running tasks and the pending isolated retry retain the ordinary two-executor floor. + assert(manager.maxNumExecutorsNeededPerResourceProfile(rpId) === 2) + draining.last.markFinished(TaskState.FINISHED, clock.getTimeMillis()) + post(SparkListenerTaskEnd(0, 0, null, Success, draining.last, new ExecutorMetrics, null)) + + // The three remaining tasks and the pending retry need no extra executor just for draining. + assert(manager.listener.pendingTasksPerResourceProfile(rpId) === 1) + assert(totalRunningTasksPerResourceProfile(manager) === 3) + assert(manager.maxNumExecutorsNeededPerResourceProfile(rpId) === 1) + + post(SparkListenerStageSubmitted(createStageInfo(1, 4))) + // Only the four new ordinary tasks need capacity outside the reserved executor. + assert(manager.listener.pendingTasksPerResourceProfile(rpId) === 5) + assert(manager.maxNumExecutorsNeededPerResourceProfile(rpId) === 2) + } + + for { + enabled <- Seq(false, true) + speculativeSuccess <- Seq(false, true) + } { + test("resubmitted shuffle demand during OOM recovery " + + s"(enabled=$enabled, speculativeSuccess=$speculativeSuccess)") { + val clock = new ManualClock(1) + val conf = createConf(0, 5, 2) + .set(config.EXECUTOR_CORES, 4) + .set(config.SCHEDULER_OOM_RETRY_ENABLED, enabled) + val rpId = DEFAULT_RESOURCE_PROFILE_ID + var reservation: Option[OomRetryReservationInfo] = None + val manager = createManager(conf, clock = clock, oomRetryReservationInfo = () => reservation) + onExecutorAdded(manager, "lost", rpManager.defaultResourceProfile) + onExecutorAdded(manager, "reserved", rpManager.defaultResourceProfile) + post(SparkListenerStageSubmitted(createStageInfo(0, 2))) + + val original = createTaskInfo(0, 0, if (speculativeSuccess) "reserved" else "lost") + post(SparkListenerTaskStart(0, 0, original)) + val completed = if (speculativeSuccess) { + post(speculativeTaskSubmitEventFromTaskIndex(0, taskIndex = 0)) + val copy = createTaskInfo(1, 0, "lost", speculative = true) + post(SparkListenerTaskStart(0, 0, copy)) + copy + } else { + original + } + completed.markFinished(TaskState.FINISHED, clock.getTimeMillis()) + post(SparkListenerTaskEnd(0, 0, null, Success, completed, new ExecutorMetrics, null)) + if (speculativeSuccess) { + original.markFinished(TaskState.KILLED, clock.getTimeMillis()) + post(SparkListenerTaskEnd(0, 0, null, TaskKilled("Another attempt succeeded"), + original, new ExecutorMetrics, null)) + } + + val retry = createTaskInfo(2, 1, "reserved") + post(SparkListenerTaskStart(0, 0, retry)) + if (enabled) { + reservation = Some(OomRetryReservationInfo(rpId, "reserved", 0, 0, 1)) + } + schedule(manager) + assert(numExecutorsTargetForDefaultProfileId(manager) === 1) + + // Losing completed output emits a second TaskEnd for the already-finished attempt. + onExecutorRemoved(manager, "lost") + post(SparkListenerTaskEnd(0, 0, null, Resubmitted, completed, new ExecutorMetrics, null)) + assert(totalRunningTasksPerResourceProfile(manager) === 1) + assert(manager.listener.pendingTasksPerResourceProfile(rpId) === 1) + clock.advance(schedulerBacklogTimeout * 1000) + schedule(manager) + // The regenerated partition needs capacity outside the still-running isolated retry. + assert(numExecutorsTargetForDefaultProfileId(manager) === (if (enabled) 2 else 1)) + + val regeneratedExecutor = if (enabled) "replacement" else "reserved" + if (enabled) { + onExecutorAdded(manager, regeneratedExecutor, rpManager.defaultResourceProfile) + } + val regenerated = createTaskInfo(3, 0, regeneratedExecutor) + post(SparkListenerTaskStart(0, 0, regenerated)) + assert(totalRunningTasksPerResourceProfile(manager) === 2) + regenerated.markFinished(TaskState.FINISHED, clock.getTimeMillis()) + post(SparkListenerTaskEnd(0, 0, null, Success, regenerated, new ExecutorMetrics, null)) + retry.markFinished(TaskState.FINISHED, clock.getTimeMillis()) + post(SparkListenerTaskEnd(0, 0, null, Success, retry, new ExecutorMetrics, null)) + reservation = None + post(SparkListenerStageCompleted(createStageInfo(0, 2))) + assert(totalRunningTasksPerResourceProfile(manager) === 0) + schedule(manager) + assert(numExecutorsTargetForDefaultProfileId(manager) === 0) + } + } + + test("OOM retry isolation counts a pending retry once across listener events") { + val conf = createConf(0, 5, 1) + .set(config.EXECUTOR_CORES, 4) + .set(config.SCHEDULER_OOM_RETRY_ENABLED, true) + val rpId = DEFAULT_RESOURCE_PROFILE_ID + var reservation = Option(OomRetryReservationInfo(rpId, "reserved", 0, 0, 0)) + val manager = createManager(conf, oomRetryReservationInfo = () => reservation) + onExecutorAdded(manager, "reserved", rpManager.defaultResourceProfile) + post(SparkListenerStageSubmitted(createStageInfo(0, 1))) + + // The scheduler snapshot can precede the retry's TaskStart event in this listener. + assert(manager.maxNumExecutorsNeededPerResourceProfile(rpId) === 1) + post(SparkListenerStageSubmitted(createStageInfo(1, 5))) + assert(manager.listener.pendingTasksPerResourceProfile(rpId) === 6) + assert(manager.maxNumExecutorsNeededPerResourceProfile(rpId) === 3) + + val retry = createTaskInfo(0, 0, "reserved") + post(SparkListenerTaskStart(0, 0, retry)) + assert(manager.listener.pendingTasksPerResourceProfile(rpId) === 5) + assert(totalRunningTasksPerResourceProfile(manager) === 1) + // Counting the retry as both pending and running would hide one ordinary executor here. + assert(manager.maxNumExecutorsNeededPerResourceProfile(rpId) === 3) + + retry.markFinished(TaskState.FINISHED, clock.getTimeMillis()) + post(SparkListenerTaskEnd(0, 0, null, Success, retry, new ExecutorMetrics, null)) + reservation = None + assert(manager.maxNumExecutorsNeededPerResourceProfile(rpId) === 2) + + // A later reservation on this executor must not also cover the finished retry. + post(SparkListenerStageSubmitted(createStageInfo(2, 1))) + reservation = Some(OomRetryReservationInfo(rpId, "reserved", 2, 0, 0)) + assert(totalRunningTasksPerResourceProfile(manager) === 0) + assert(manager.maxNumExecutorsNeededPerResourceProfile(rpId) === 3) + + reservation = None + post(SparkListenerStageCompleted(createStageInfo(0, 1))) + post(SparkListenerStageCompleted(createStageInfo(1, 5))) + post(SparkListenerStageCompleted(createStageInfo(2, 1))) + assert(manager.maxNumExecutorsNeededPerResourceProfile(rpId) === 0) + } + + test("OOM retry isolation does not subtract late tasks from a canceled stage") { + val clock = new ManualClock(1) + val conf = createConf(0, 5, 1) + .set(config.EXECUTOR_CORES, 4) + .set(config.SCHEDULER_OOM_RETRY_ENABLED, true) + val rpId = DEFAULT_RESOURCE_PROFILE_ID + var reservation = Option(OomRetryReservationInfo(rpId, "reserved", 0, 0, 0)) + val manager = createManager(conf, clock = clock, oomRetryReservationInfo = () => reservation) + onExecutorAdded(manager, "reserved", rpManager.defaultResourceProfile) + post(SparkListenerStageSubmitted(createStageInfo(0, 1))) + post(SparkListenerStageCompleted(createStageInfo(0, 1))) + post(SparkListenerStageSubmitted(createStageInfo(1, 1))) + assert(manager.maxNumExecutorsNeededPerResourceProfile(rpId) === 2) + + // Cancellation can remove the stage's resource-profile association before a queued start. + val retry = createTaskInfo(0, 0, "reserved") + post(SparkListenerTaskStart(0, 0, retry)) + assert(manager.listener.pendingTasksPerResourceProfile(rpId) === 1) + assert(totalRunningTasksPerResourceProfile(manager) === 0) + assert(manager.maxNumExecutorsNeededPerResourceProfile(rpId) === 2) + clock.advance(schedulerBacklogTimeout * 1000) + schedule(manager) + assert(numExecutorsTargetForDefaultProfileId(manager) === 2) + + retry.markFinished(TaskState.KILLED, clock.getTimeMillis()) + post(SparkListenerTaskEnd(0, 0, null, TaskKilled("Stage cancelled"), + retry, new ExecutorMetrics, null)) + reservation = None + schedule(manager) + assert(numExecutorsTargetForDefaultProfileId(manager) === 1) + post(SparkListenerStageCompleted(createStageInfo(1, 1))) + schedule(manager) + assert(numExecutorsTargetForDefaultProfileId(manager) === 0) + } + + test("OOM retry isolation preserves the extra executor for ordinary speculation") { + val clock = new ManualClock(1) + val conf = createConf(0, 5, 2) + .set(config.EXECUTOR_CORES, 4) + .set(config.SCHEDULER_OOM_RETRY_ENABLED, true) + val rpId = DEFAULT_RESOURCE_PROFILE_ID + var reservation = Option(OomRetryReservationInfo(rpId, "reserved", 0, 0, 0)) + val manager = createManager(conf, clock = clock, oomRetryReservationInfo = () => reservation) + onExecutorAdded(manager, "reserved", rpManager.defaultResourceProfile) + onExecutorAdded(manager, "ordinary", rpManager.defaultResourceProfile) + post(SparkListenerStageSubmitted(createStageInfo(0, 1))) + post(SparkListenerTaskStart(0, 0, createTaskInfo(0, 0, "reserved"))) + post(SparkListenerStageSubmitted(createStageInfo(1, 1))) + post(SparkListenerTaskStart(1, 0, createTaskInfo(1, 0, "ordinary"))) + post(speculativeTaskSubmitEventFromTaskIndex(1, taskIndex = 0)) + + assert(totalRunningTasksPerResourceProfile(manager) === 2) + assert(manager.listener.pendingSpeculativeTasksPerResourceProfile(rpId) === 1) + assert(manager.maxNumExecutorsNeededPerResourceProfile(rpId) === 3) + clock.advance(schedulerBacklogTimeout * 1000) + schedule(manager) + assert(numExecutorsTargetForDefaultProfileId(manager) === 3) + + reservation = None + schedule(manager) + assert(manager.maxNumExecutorsNeededPerResourceProfile(rpId) === 2) + assert(numExecutorsTargetForDefaultProfileId(manager) === 2) + } + + test("OOM retry isolation separates speculation from an original on the reserved executor") { + val clock = new ManualClock(1) + val conf = createConf(0, 5, 2) + .set(config.EXECUTOR_CORES, 4) + .set(config.SCHEDULER_OOM_RETRY_ENABLED, true) + val rpId = DEFAULT_RESOURCE_PROFILE_ID + var reservation = Option(OomRetryReservationInfo(rpId, "reserved", 0, 0, 1)) + val manager = createManager(conf, clock = clock, oomRetryReservationInfo = () => reservation) + onExecutorAdded(manager, "reserved", rpManager.defaultResourceProfile) + post(SparkListenerStageSubmitted(createStageInfo(0, 2))) + val original = createTaskInfo(0, 0, "reserved") + post(SparkListenerTaskStart(0, 0, original)) + post(speculativeTaskSubmitEventFromTaskIndex(0, taskIndex = 0)) + + assert(manager.listener.pendingTasksPerResourceProfile(rpId) === 1) + assert(manager.listener.pendingSpeculativeTasksPerResourceProfile(rpId) === 1) + // The ordinary executor is already separate from the draining original's executor. + assert(manager.maxNumExecutorsNeededPerResourceProfile(rpId) === 2) + clock.advance(schedulerBacklogTimeout * 1000) + schedule(manager) + assert(numExecutorsTargetForDefaultProfileId(manager) === 2) + + onExecutorAdded(manager, "ordinary", rpManager.defaultResourceProfile) + val speculative = createTaskInfo(1, 0, "ordinary", speculative = true) + post(SparkListenerTaskStart(0, 0, speculative)) + original.markFinished(TaskState.FINISHED, clock.getTimeMillis()) + post(SparkListenerTaskEnd(0, 0, null, Success, original, new ExecutorMetrics, null)) + assert(manager.listener.pendingSpeculativeTasksPerResourceProfile(rpId) === 0) + // The finished original must no longer cover the copy still running elsewhere. + assert(manager.maxNumExecutorsNeededPerResourceProfile(rpId) === 2) + + speculative.markFinished(TaskState.FINISHED, clock.getTimeMillis()) + post(SparkListenerTaskEnd(0, 0, null, Success, speculative, new ExecutorMetrics, null)) + assert(manager.maxNumExecutorsNeededPerResourceProfile(rpId) === 1) + val retry = createTaskInfo(2, 1, "reserved") + post(SparkListenerTaskStart(0, 0, retry)) + retry.markFinished(TaskState.FINISHED, clock.getTimeMillis()) + post(SparkListenerTaskEnd(0, 0, null, Success, retry, new ExecutorMetrics, null)) + reservation = None + schedule(manager) + assert(numExecutorsTargetForDefaultProfileId(manager) === 0) + } + + Seq(("disabled", false, 4), ("single-slot", true, 1)).foreach { + case (name, enabled, cores) => + test(s"OOM retry isolation leaves $name allocation unchanged") { + val conf = createConf(0, 5, 1) + .set(config.EXECUTOR_CORES, cores) + .set(config.SCHEDULER_OOM_RETRY_ENABLED, enabled) + .set(config.DYN_ALLOCATION_EXECUTOR_ALLOCATION_RATIO, 0.1) + val rpId = DEFAULT_RESOURCE_PROFILE_ID + var reservation = Option(OomRetryReservationInfo(rpId, "reserved", 0, 0, 0)) + val manager = createManager(conf, oomRetryReservationInfo = () => reservation) + onExecutorAdded(manager, "reserved", rpManager.defaultResourceProfile) + post(SparkListenerStageSubmitted(createStageInfo(0, 4))) + post(SparkListenerTaskStart(0, 0, createTaskInfo(0, 0, "reserved"))) + + assert(manager.maxNumExecutorsNeededPerResourceProfile(rpId) === 1) + reservation = None + assert(manager.maxNumExecutorsNeededPerResourceProfile(rpId) === 1) + } + } + + test("OOM retry isolation only changes demand for the reserved resource profile") { + val conf = createConf(0, 5, 1) + .set(config.EXECUTOR_CORES, 4) + .set(config.SCHEDULER_OOM_RETRY_ENABLED, true) + var reservation: Option[OomRetryReservationInfo] = None + val manager = createManager(conf, oomRetryReservationInfo = () => reservation) + val otherProfile = new ResourceProfileBuilder() + .require(new ExecutorResourceRequests().cores(4).resource("gpu", 4)) + .require(new TaskResourceRequests().cpus(1).resource("gpu", 1)) + .build() + rpManager.addResourceProfile(otherProfile) + val rpId = DEFAULT_RESOURCE_PROFILE_ID + onExecutorAdded(manager, "default", rpManager.defaultResourceProfile) + onExecutorAdded(manager, "other", otherProfile) + post(SparkListenerStageSubmitted(createStageInfo(0, 4))) + post(SparkListenerTaskStart(0, 0, createTaskInfo(0, 0, "default"))) + post(SparkListenerStageSubmitted(createStageInfo(1, 4, rp = otherProfile))) + post(SparkListenerTaskStart(1, 0, createTaskInfo(1, 0, "other"))) + + reservation = Some(OomRetryReservationInfo(otherProfile.id, "other", 1, 0, 0)) + assert(manager.maxNumExecutorsNeededPerResourceProfile(rpId) === 1) + assert(manager.maxNumExecutorsNeededPerResourceProfile(otherProfile.id) === 2) + reservation = Some(OomRetryReservationInfo(rpId, "default", 0, 0, 0)) + assert(manager.maxNumExecutorsNeededPerResourceProfile(rpId) === 2) + assert(manager.maxNumExecutorsNeededPerResourceProfile(otherProfile.id) === 1) + } + test("SPARK-31418: one stage being unschedulable") { val clock = new ManualClock() val conf = createConf(0, 5, 0).set(config.EXECUTOR_CORES, 2) @@ -1996,12 +2409,15 @@ class ExecutorAllocationManagerSuite extends SparkFunSuite { private def createManager( conf: SparkConf, - clock: Clock = new SystemClock()): ExecutorAllocationManager = { + clock: Clock = new SystemClock(), + oomRetryReservationInfo: () => Option[OomRetryReservationInfo] = () => None): + ExecutorAllocationManager = { ResourceProfile.reInitDefaultProfile(conf) rpManager = new ResourceProfileManager(conf, listenerBus) val manager = new ExecutorAllocationManager(client, listenerBus, conf, clock = clock, - resourceProfileManager = rpManager, reliableShuffleStorage = false) + resourceProfileManager = rpManager, reliableShuffleStorage = false, + oomRetryReservationInfo = oomRetryReservationInfo) managers += manager manager.start() manager diff --git a/core/src/test/scala/org/apache/spark/FailureSuite.scala b/core/src/test/scala/org/apache/spark/FailureSuite.scala index 8b75c3a0ba653..3c5616e3d0c28 100644 --- a/core/src/test/scala/org/apache/spark/FailureSuite.scala +++ b/core/src/test/scala/org/apache/spark/FailureSuite.scala @@ -18,11 +18,20 @@ package org.apache.spark import java.io.{IOException, NotSerializableException, ObjectInputStream} +import java.util.Collections +import java.util.concurrent.{CountDownLatch, TimeUnit} + +import scala.collection.mutable.ArrayBuffer +import scala.concurrent.{ExecutionContext, Future} +import scala.concurrent.duration._ + +import org.scalatest.concurrent.Eventually import org.apache.spark.internal.config.UNSAFE_EXCEPTION_ON_MEMORY_LEAK -import org.apache.spark.memory.TestMemoryConsumer +import org.apache.spark.memory.{SparkOutOfMemoryError, TestMemoryConsumer} +import org.apache.spark.scheduler.TaskSchedulerImpl import org.apache.spark.storage.StorageLevel -import org.apache.spark.util.NonSerializable +import org.apache.spark.util.{NonSerializable, ThreadUtils, Utils} // Common state shared by FailureSuite-launched tasks. We use a global object // for this because any local variables used in the task closures will rightfully @@ -30,16 +39,26 @@ import org.apache.spark.util.NonSerializable object FailureSuiteState { var tasksRun = 0 var tasksFailed = 0 + @volatile var oomBlockerStarted = new CountDownLatch(1) + @volatile var secondOom = new CountDownLatch(1) + @volatile var oomRetryStarted = new CountDownLatch(1) + @volatile var ordinaryTaskStarted = new CountDownLatch(1) + @volatile var releaseOomBlocker = new CountDownLatch(1) def clear(): Unit = { synchronized { tasksRun = 0 tasksFailed = 0 + oomBlockerStarted = new CountDownLatch(1) + secondOom = new CountDownLatch(1) + oomRetryStarted = new CountDownLatch(1) + ordinaryTaskStarted = new CountDownLatch(1) + releaseOomBlocker = new CountDownLatch(1) } } } -class FailureSuite extends SparkFunSuite with LocalSparkContext { +class FailureSuite extends SparkFunSuite with LocalSparkContext with Eventually { // Run a 3-task map job in which task 1 deterministically fails once, and check // whether the job completes successfully and we ran 4 tasks in total. @@ -222,6 +241,231 @@ class FailureSuite extends SparkFunSuite with LocalSparkContext { FailureSuiteState.clear() } + test("OOM retries preserve task CPUs and failure limits through the local backend") { + sc = new SparkContext(new SparkConf() + .setMaster("local[2,4]") + .setAppName("OOM retry isolation") + .set("spark.scheduler.oomRetry.enabled", "true")) + val attempts = sc.parallelize(Seq(0), 1).mapPartitions { _ => + val context = TaskContext.get() + if (context.attemptNumber() < 2) { + // scalastyle:off throwerror + throw new SparkOutOfMemoryError("_LEGACY_ERROR_USER_RAISED_EXCEPTION", + Collections.singletonMap("errorMessage", "execution memory")) + // scalastyle:on throwerror + } + Iterator((context.attemptNumber(), context.cpus())) + }.collect() + assert(attempts.toSeq == Seq((2, 1))) + + val failure = intercept[SparkException] { + sc.parallelize(Seq(0), 1).foreach { _ => + // scalastyle:off throwerror + throw new SparkOutOfMemoryError("_LEGACY_ERROR_USER_RAISED_EXCEPTION", + Collections.singletonMap("errorMessage", "persistent execution memory failure")) + // scalastyle:on throwerror + } + } + assert(failure.getMessage.contains("failed 4 times")) + } + + test("OOM isolation timeout wakes the local backend while another task is running") { + FailureSuiteState.clear() + sc = new SparkContext(new SparkConf() + .setMaster("local[2,4]") + .setAppName("OOM retry isolation timeout") + .set("spark.scheduler.oomRetry.enabled", "true") + .set("spark.scheduler.oomRetry.isolationTimeout", "1s")) + val pool = ThreadUtils.newDaemonSingleThreadExecutor("oom-retry-timeout-test") + val executionContext = ExecutionContext.fromExecutorService(pool) + val context = sc + val job = Future { + context.parallelize(Seq(0, 1), 2).mapPartitionsWithIndex { (index, _) => + val attempt = TaskContext.get().attemptNumber() + if (index == 0) { + FailureSuiteState.oomBlockerStarted.countDown() + require(FailureSuiteState.releaseOomBlocker.await(60, TimeUnit.SECONDS), + "The retry did not release the blocker") + } else { + require(FailureSuiteState.oomBlockerStarted.await(30, TimeUnit.SECONDS), + "The blocker did not start") + if (attempt < 2) { + if (attempt == 1) { + FailureSuiteState.secondOom.countDown() + } + // scalastyle:off throwerror + throw new SparkOutOfMemoryError("_LEGACY_ERROR_USER_RAISED_EXCEPTION", + Collections.singletonMap("errorMessage", "execution memory")) + // scalastyle:on throwerror + } + FailureSuiteState.oomRetryStarted.countDown() + FailureSuiteState.releaseOomBlocker.countDown() + } + Iterator((index, attempt)) + }.collect() + }(executionContext) + try { + assert(FailureSuiteState.secondOom.await(30, TimeUnit.SECONDS), + "The task did not fail twice with OOM") + // The blocker stays running, so only the deadline can make the retry runnable again. + assert(FailureSuiteState.oomRetryStarted.await(10, TimeUnit.SECONDS), + "The isolation deadline passed without waking the local backend") + assert(ThreadUtils.awaitResult(job, 30.seconds).sorted.toSeq == Seq((0, 0), (1, 2))) + } finally { + FailureSuiteState.releaseOomBlocker.countDown() + try { + context.cancelAllJobs() + ThreadUtils.awaitReady(job, 30.seconds) + } finally { + pool.shutdownNow() + assert(pool.awaitTermination(30, TimeUnit.SECONDS)) + FailureSuiteState.clear() + } + } + } + + test("cancelling a pending OOM reservation wakes queued work through the local backend") { + FailureSuiteState.clear() + sc = new SparkContext(new SparkConf() + .setMaster("local[2,4]") + .setAppName("OOM retry cancellation") + .set("spark.scheduler.oomRetry.enabled", "true") + .set("spark.scheduler.oomRetry.isolationTimeout", "60s")) + val pool = ThreadUtils.newDaemonFixedThreadPool(3, "oom-retry-cancellation-test") + val executionContext = ExecutionContext.fromExecutorService(pool) + val context = sc + val scheduler = context.taskScheduler.asInstanceOf[TaskSchedulerImpl] + val jobs = ArrayBuffer.empty[Future[_]] + try { + val blocker = Future { + context.setJobGroup("oom-blocker", "occupy one core", interruptOnCancel = true) + context.parallelize(Seq(0), 1).map { _ => + FailureSuiteState.oomBlockerStarted.countDown() + require(FailureSuiteState.releaseOomBlocker.await(60, TimeUnit.SECONDS), + "The queued task did not release the blocker") + 0 + }.collect() + }(executionContext) + jobs += blocker + assert(FailureSuiteState.oomBlockerStarted.await(10, TimeUnit.SECONDS), + "The blocker did not start") + + val oom = Future { + context.setJobGroup("oom-retry", "reserve the executor", interruptOnCancel = true) + context.parallelize(Seq(1), 1).map { _ => + val attempt = TaskContext.get().attemptNumber() + if (attempt < 2) { + if (attempt == 1) { + FailureSuiteState.secondOom.countDown() + } + // scalastyle:off throwerror + throw new SparkOutOfMemoryError("_LEGACY_ERROR_USER_RAISED_EXCEPTION", + Collections.singletonMap("errorMessage", "execution memory")) + // scalastyle:on throwerror + } + FailureSuiteState.oomRetryStarted.countDown() + 1 + }.collect() + }(executionContext) + jobs += oom + assert(FailureSuiteState.secondOom.await(10, TimeUnit.SECONDS), + "The task did not fail twice with OOM") + eventually(timeout(10.seconds)) { + scheduler.synchronized { + val manager = scheduler.rootPool.getSortedTaskSetQueue.find { + _.taskSet.properties.getProperty("spark.jobGroup.id") == "oom-retry" + }.get + assert(manager.taskAttempts.head.size == 2) + assert(manager.taskAttempts.head.forall(_.failed)) + assert(manager.runningTasks == 0) + } + } + + val ordinary = Future { + context.setJobGroup("ordinary", "queued ordinary task", interruptOnCancel = true) + context.parallelize(Seq(2), 1).map { _ => + FailureSuiteState.ordinaryTaskStarted.countDown() + FailureSuiteState.releaseOomBlocker.countDown() + 2 + }.collect() + }(executionContext) + jobs += ordinary + eventually(timeout(10.seconds)) { + assert(scheduler.synchronized { + scheduler.rootPool.getSortedTaskSetQueue.exists { manager => + manager.taskSet.properties.getProperty("spark.jobGroup.id") == "ordinary" && + manager.runningTasks == 0 + } + }) + } + assert(!FailureSuiteState.ordinaryTaskStarted.await(200, TimeUnit.MILLISECONDS), + "The reservation did not block ordinary work on the free core") + + // No attempt of this job is running, so cancellation cannot produce a task completion + // offer. The queued job must start before either the blocker or isolation deadline expires. + context.cancelJobGroup("oom-retry") + ThreadUtils.awaitReady(oom, 10.seconds) + assert(oom.value.get.isFailure) + assert(FailureSuiteState.ordinaryTaskStarted.await(10, TimeUnit.SECONDS), + "Cancelling the reservation did not wake the local backend") + assert(FailureSuiteState.oomRetryStarted.getCount == 1) + assert(ThreadUtils.awaitResult(ordinary, 10.seconds).toSeq == Seq(2)) + assert(ThreadUtils.awaitResult(blocker, 10.seconds).toSeq == Seq(0)) + } finally { + FailureSuiteState.releaseOomBlocker.countDown() + try { + context.cancelAllJobs() + jobs.foreach(job => ThreadUtils.awaitReady(job, 10.seconds)) + } finally { + pool.shutdownNow() + assert(pool.awaitTermination(10, TimeUnit.SECONDS)) + FailureSuiteState.clear() + } + } + } + + test("ExceptionFailure identifies typed OOM causes without matching exception text") { + val errors = Seq( + new OutOfMemoryError("heap"), + new SparkOutOfMemoryError("_LEGACY_ERROR_USER_RAISED_EXCEPTION", + Collections.singletonMap("errorMessage", "execution memory"))) + for (error <- errors; preserveCause <- Seq(true, false)) { + assert(new ExceptionFailure(error, Nil, preserveCause).isOutOfMemoryError) + val wrapped = new RuntimeException("spill failed", error) + assert(new ExceptionFailure(wrapped, Nil, preserveCause).isOutOfMemoryError) + assert(ExceptionFailure(error.getClass.getName, error.getMessage, error.getStackTrace, + Utils.exceptionString(error), None).isOutOfMemoryError) + } + val ordinary = new RuntimeException("java.lang.OutOfMemoryError: OOMKilled") + assert(!new ExceptionFailure(ordinary, Nil).isOutOfMemoryError) + + val first = new RuntimeException("first") + val second = new RuntimeException("second", first) + first.initCause(second) + assert(!new ExceptionFailure(first, Nil).isOutOfMemoryError) + } + + test("ExceptionFailure preserves wrapped OOM classification without a serializable cause") { + val error = new NonSerializableUserException + error.initCause(new OutOfMemoryError("heap")) + intercept[NotSerializableException] { + Utils.serialize(new ExceptionFailure(error, Nil)) + } + val fallback = new ExceptionFailure(error, Nil, preserveCause = false) + val restored = Utils.deserialize[ExceptionFailure](Utils.serialize(fallback)) + assert(restored.exception.isEmpty) + assert(restored.isOutOfMemoryError) + } + + test("ExceptionFailure preserves wrapped OOM classification if its cause cannot deserialize") { + val error = new NonDeserializableUserException + error.initCause(new OutOfMemoryError("heap")) + val restored = Utils.deserialize[ExceptionFailure]( + Utils.serialize(new ExceptionFailure(error, Nil))) + assert(restored.exception.isEmpty) + assert(restored.isOutOfMemoryError) + } + // Run a 3-task map stage where one task fails once. test("failure in tasks in a submitMapStage") { sc = new SparkContext("local[1,2]", "test") diff --git a/core/src/test/scala/org/apache/spark/scheduler/CoarseGrainedSchedulerBackendSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/CoarseGrainedSchedulerBackendSuite.scala index ecdd02585a123..eabfd244fd0f0 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/CoarseGrainedSchedulerBackendSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/CoarseGrainedSchedulerBackendSuite.scala @@ -35,6 +35,7 @@ import org.scalatestplus.mockito.MockitoSugar._ import org.apache.spark._ import org.apache.spark.TestUtils.createTempScriptWithExpectedOutput +import org.apache.spark.errors.SparkCoreErrors import org.apache.spark.internal.config._ import org.apache.spark.internal.config.Network.RPC_MESSAGE_MAX_SIZE import org.apache.spark.rdd.RDD @@ -64,10 +65,90 @@ class CoarseGrainedSchedulerBackendSuite extends SparkFunSuite with LocalSparkCo larger.collect() } assert(thrown.getMessage.contains("using broadcast variables for large values")) + val scheduler = sc.taskScheduler.asInstanceOf[TaskSchedulerImpl] + eventually(timeout(5.seconds)) { + assert(scheduler.taskIdToTaskSetManager.isEmpty) + assert(scheduler.runningTasksByExecutors.values.forall(_ == 0)) + } val smaller = sc.parallelize(1 to 4).collect() assert(smaller.length === 4) } + test("RPC size rejection clears an isolated OOM retry without refunding resources") { + val conf = new SparkConf() + .set(RPC_MESSAGE_MAX_SIZE, 1) + .set(SCHEDULER_OOM_RETRY_ENABLED, true) + .set(EXECUTOR_CORES, 2) + .set(EXECUTOR_GPU_ID.amountConf, "1") + .set(TASK_GPU_ID.amountConf, "1") + val backend = createDecommissionBackend(conf) + val scheduler = backend.taskScheduler + val resources = Map(GPU -> new ResourceInformation(GPU, Array("0"))) + val executor = registerDecommissionExecutor(backend, "1", 2, resources) + val taskSet = FakeTask.createTaskSet(1) + scheduler.submitTasks(taskSet) + val manager = scheduler.taskSetManagerForAttempt(0, 0).get + val frameSize = RpcUtils.maxMessageSizeBytes(sc.conf) + + (0 until 2).foreach { attempt => + backend.driverEndpoint.send(ReviveOffers) + val task = executor.nextTask() + assert(task.attemptNumber === attempt) + assert(backend.getExecutorAvailableCpus("1").contains(1)) + assert(backend.getExecutorAvailableResources("1")(GPU).availableAddrs.isEmpty) + if (attempt == 1) { + // Properties are encoded into the launch message, so only the third attempt is too big. + taskSet.tasks(0).localProperties.setProperty("large", "x" * (2 * frameSize)) + } + val failure = new ExceptionFailure(SparkCoreErrors.outOfMemoryError(1, 0, ""), Nil) + val serializedFailure = sc.env.closureSerializer.newInstance().serialize(failure) + backend.driverEndpoint.send(StatusUpdate( + task.executorId, task.taskId, TaskState.FAILED, new SerializableBuffer(serializedFailure), + task.cpus, task.resources)) + flushDecommissionBackend(backend) + eventually(timeout(5.seconds)) { + scheduler.synchronized { + assert(manager.taskInfos(task.taskId).finished) + } + } + } + + scheduler.synchronized { + assert(manager.oomRetryNeedsIsolation(0)) + } + backend.driverEndpoint.send(ReviveOffers) + flushDecommissionBackend(backend) + eventually(timeout(5.seconds)) { + scheduler.synchronized { + assert(manager.isZombie) + assert(manager.taskAttempts(0).size === 3) + assert(manager.taskAttempts(0).head.killed) + assert(manager.runningTasks === 0) + assert(scheduler.taskSetManagerForAttempt(0, 0).isEmpty) + assert(scheduler.taskIdToTaskSetManager.isEmpty) + assert(scheduler.oomRetryReservationInfo.isEmpty) + assert(!scheduler.isExecutorBusy("1")) + } + } + assert(executor.launchedTasks.isEmpty) + assert(backend.getExecutorAvailableCpus("1").contains(2)) + assert(backend.getExecutorAvailableResources("1")(GPU).availableAddrs === Array("0")) + + scheduler.submitTasks(FakeTask.createTaskSet(1, stageId = 1, stageAttemptId = 0)) + val ordinaryManager = scheduler.taskSetManagerForAttempt(1, 0).get + backend.driverEndpoint.send(ReviveOffers) + val ordinary = executor.nextTask() + completeDecommissionTestTask(backend, ordinary) + eventually(timeout(5.seconds)) { + scheduler.synchronized { + assert(ordinaryManager.taskInfos(ordinary.taskId).successful) + } + } + assert(!scheduler.isExecutorBusy("1")) + assert(backend.getExecutorAvailableCpus("1").contains(2)) + assert(backend.getExecutorAvailableResources("1")(GPU).availableAddrs === Array("0")) + } + test("compute max number of concurrent tasks can be launched") { val conf = new SparkConf() .setMaster("local-cluster[4, 3, 1024]") @@ -1034,12 +1115,13 @@ class CoarseGrainedSchedulerBackendSuite extends SparkFunSuite with LocalSparkCo private def registerDecommissionExecutor( backend: DecommissionTestSchedulerBackend, executorId: String, - cores: Int = 1) + cores: Int = 1, + resources: Map[String, ResourceInformation] = Map.empty) : DecommissionTestExecutorRpcEndpointRef = { val executor = new DecommissionTestExecutorRpcEndpointRef(sc.conf, executorId) assert(backend.driverEndpoint.askSync[Boolean]( RegisterExecutor(executorId, executor, "localhost", cores, Map.empty, Map.empty, - Map.empty, ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID))) + resources, ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID))) backend.driverEndpoint.send(LaunchedExecutor(executorId)) flushDecommissionBackend(backend) executor diff --git a/core/src/test/scala/org/apache/spark/scheduler/ExecutorResourcesAmountsSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/ExecutorResourcesAmountsSuite.scala index 1c5cb041ad6c0..dfbe50314dfdd 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/ExecutorResourcesAmountsSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/ExecutorResourcesAmountsSuite.scala @@ -39,6 +39,40 @@ class ExecutorResourcesAmountsSuite extends SparkFunSuite with ExecutorResourceU k -> ResourceAmountUtils.toInternalResource(v) } } + test("available task slots use current amounts without consuming the offer") { + val resources = new ExecutorResourcesAmounts( + toInternalResourceMap(Map(GPU -> Map("0" -> 1.0, "1" -> 1.0)))) + val profile = new ResourceProfileBuilder() + .require(new TaskResourceRequests().resource(GPU, 1)).build() + assert(resources.availableTaskSlots(profile, 1) === 1) + val assigned = resources.assignAddressesCustomResources(profile).get + resources.acquire(assigned) + val before = resources.availableResources + assert(resources.availableTaskSlots(profile, 8) === 1) + assert(resources.availableResources === before) + resources.release(assigned) + assert(resources.availableTaskSlots(profile, 8) === 2) + } + + test("available task slots respect fractional resource fragmentation") { + val resources = new ExecutorResourcesAmounts( + toInternalResourceMap(Map(GPU -> Map("0" -> 0.25, "1" -> 0.75)))) + val fractional = new ResourceProfileBuilder() + .require(new TaskResourceRequests().resource(GPU, 0.5)).build() + val whole = new ResourceProfileBuilder() + .require(new TaskResourceRequests().resource(GPU, 1)).build() + assert(resources.availableTaskSlots(fractional, 8) === 1) + assert(resources.availableTaskSlots(whole, 8) === 0) + } + + test("available task slots check every requested resource") { + val resources = new ExecutorResourcesAmounts(toInternalResourceMap(Map( + GPU -> Map("0" -> 1.0, "1" -> 1.0), "fpga" -> Map("0" -> 1.0)))) + val profile = new ResourceProfileBuilder() + .require(new TaskResourceRequests().resource(GPU, 0.5).resource("fpga", 1)).build() + assert(resources.availableTaskSlots(profile, 8) === 1) + } + test("assign to rp without task resources requirement") { val executorsInfo = Map( "gpu" -> new ExecutorResourceInfo("gpu", Seq("2", "4", "6")), diff --git a/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala index 810a0b7c483b5..135582eaa1d6f 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala @@ -28,12 +28,13 @@ import scala.concurrent.duration._ import scala.language.implicitConversions import scala.language.reflectiveCalls -import org.mockito.ArgumentMatchers.{any, anyInt, anyString, eq => meq} -import org.mockito.Mockito.{atLeast, atMost, never, spy, times, verify, when} +import org.mockito.ArgumentMatchers.{any, anyBoolean, anyInt, anyLong, anyString, eq => meq} +import org.mockito.Mockito.{atLeast, atMost, clearInvocations, never, spy, times, verify, when} import org.scalatest.concurrent.Eventually import org.scalatestplus.mockito.MockitoSugar import org.apache.spark._ +import org.apache.spark.executor.TaskMetrics import org.apache.spark.internal.config import org.apache.spark.resource.{CpuAmount, ExecutorResourceRequests, ResourceAmountUtils, ResourceProfile, TaskResourceProfile, TaskResourceRequests} import org.apache.spark.resource.ResourceAmountUtils.ONE_ENTIRE_RESOURCE @@ -208,6 +209,738 @@ class TaskSchedulerImplSuite extends SparkFunSuite with LocalSparkContext assert(!failedTaskSet) } + private def setupOomRetryScheduler( + clock: Clock, + confs: (String, String)*): TaskSchedulerImpl = { + val conf = new SparkConf().setMaster("local[8]").setAppName("TaskSchedulerImplSuite") + .set(config.SCHEDULER_OOM_RETRY_ENABLED, true) + confs.foreach { case (k, v) => conf.set(k, v) } + sc = new SparkContext(conf) + taskScheduler = new TaskSchedulerImpl(sc, sc.conf.get(config.TASK_MAX_FAILURES), + clock = clock) { + override def shuffleOffers(offers: IndexedSeq[WorkerOffer]): IndexedSeq[WorkerOffer] = offers + } + setupHelper() + } + + private def oomReservationInfo( + manager: TaskSetManager, + executorId: String): OomRetryReservationInfo = { + val taskSet = manager.taskSet + OomRetryReservationInfo( + taskSet.resourceProfileId, executorId, taskSet.stageId, taskSet.stageAttemptId, 0) + } + + private def failTaskWithOom(task: TaskDescription): Unit = taskScheduler.synchronized { + val manager = taskScheduler.taskIdToTaskSetManager.get(task.taskId) + // Keep the synchronous failure handler ahead of TaskResultGetter's asynchronous callback. + failTask(task.taskId, TaskState.FAILED, + new ExceptionFailure(new OutOfMemoryError("test OOM"), Nil), manager) + } + + private def finishOomTestTask(task: TaskDescription): Unit = { + val manager = taskScheduler.taskIdToTaskSetManager.get(task.taskId) + val value = sc.env.serializer.newInstance().serialize(0) + val result = new DirectTaskResult[Int](value, Seq.empty, Array.empty[Long]) + taskScheduler.statusUpdate(task.taskId, TaskState.FINISHED, + sc.env.closureSerializer.newInstance().serialize(result)) + eventually(timeout(10.seconds)) { + assert(manager.taskInfos(task.taskId).finished) + assert(manager.successful(task.index)) + } + } + + private def prepareOomRetries(numTasks: Int = 1, stageId: Int = 1): TaskSetManager = { + taskScheduler.submitTasks(FakeTask.createTaskSet(numTasks, stageId, stageAttemptId = 0)) + val manager = taskScheduler.taskSetManagerForAttempt(stageId, 0).get + (0 until 2).foreach { _ => + val tasks = taskScheduler.resourceOffers( + IndexedSeq(WorkerOffer("origin", "origin-host", numTasks))).flatten + assert(tasks.size === numTasks) + assert(tasks.forall(task => + taskScheduler.taskIdToTaskSetManager.get(task.taskId) eq manager)) + tasks.foreach(failTaskWithOom) + } + manager + } + + test("OOM retries prefer idle executors before ordinary work") { + val scheduler = setupOomRetryScheduler(new ManualClock(1), + config.DYN_ALLOCATION_ENABLED.key -> "true") + scheduler.submitTasks(FakeTask.createTaskSet(1, stageId = 0, stageAttemptId = 0)) + scheduler.resourceOffers(IndexedSeq(WorkerOffer("a", "host-a", 1))) + scheduler.submitTasks(FakeTask.createTaskSet(1, stageId = 2, stageAttemptId = 0)) + val original = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("a", "host-a", 1))).flatten.head + val manager = scheduler.taskSetManagerForAttempt(2, 0).get + failTaskWithOom(original) + + // This TaskSet precedes the retry in FIFO order and could consume every offered slot. + scheduler.submitTasks(FakeTask.createTaskSet(7, stageId = 1, stageAttemptId = 0)) + val tasks = scheduler.resourceOffers(IndexedSeq( + WorkerOffer("a", "host-a", 3), WorkerOffer("b", "host-b", 4))).flatten + val retries = tasks.filter(task => + scheduler.taskIdToTaskSetManager.get(task.taskId) eq manager) + + assert(tasks.size === 7) + assert(retries.map(_.executorId) === Seq("b")) + assert(retries.head.attemptNumber === 1) + assert(retries.head.cpus === 1) + // A first OOM changes placement only; ordinary tasks can still share its executor. + assert(tasks.count(_.executorId == "b") === 4) + } + + test("a first OOM retry uses busy capacity when no executor is idle") { + val scheduler = setupOomRetryScheduler(new ManualClock(1)) + scheduler.submitTasks(FakeTask.createTaskSet(2)) + val originals = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("a", "host-a", 2))).flatten + failTaskWithOom(originals.head) + + val tasks = scheduler.resourceOffers(IndexedSeq(WorkerOffer("a", "host-a", 2))).flatten + assert(tasks.map(_.index) === Seq(originals.head.index)) + assert(tasks.head.executorId === "a") + assert(tasks.head.cpus === 1) + } + + test("OOM recovery preserves fractional CPU accounting before and after isolation") { + val scheduler = setupOomRetryScheduler(new ManualClock(1), + config.CPUS_PER_TASK.key -> "0.2") + val taskCpus = BigDecimal("0.2") + scheduler.submitTasks(FakeTask.createTaskSet(1, stageId = 1, stageAttemptId = 0)) + val original = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("origin", "host-origin", 1))).flatten.head + val manager = scheduler.taskSetManagerForAttempt(1, 0).get + failTaskWithOom(original) + + scheduler.submitTasks(FakeTask.createTaskSet(10, stageId = 0, stageAttemptId = 0)) + val first = scheduler.resourceOffers(IndexedSeq(WorkerOffer("a", "host-a", 1))).flatten + val retry = first.find(task => + scheduler.taskIdToTaskSetManager.get(task.taskId) eq manager).get + assert(first.size === 5) + assert(first.forall(_.cpus == taskCpus)) + assert(first.map(_.cpus).sum === BigDecimal(1)) + failTaskWithOom(retry) + + val isolated = scheduler.resourceOffers(IndexedSeq( + WorkerOffer("a", "host-a", taskCpus), WorkerOffer("b", "host-b", 1))).flatten + val secondRetry = isolated.find(task => + scheduler.taskIdToTaskSetManager.get(task.taskId) eq manager).get + assert(secondRetry.executorId === "b") + assert(secondRetry.cpus === taskCpus) + assert(isolated.count(_.executorId == "b") === 1) + assert(isolated.count(_.executorId == "a") === 1) + assert(scheduler.resourceOffers( + IndexedSeq(WorkerOffer("b", "host-b", BigDecimal("0.8")))).flatten.isEmpty) + + finishOomTestTask(secondRetry) + val remaining = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("b", "host-b", 1))).flatten + assert(remaining.size === 5) + assert(remaining.map(_.cpus).sum === BigDecimal(1)) + } + + test("repeated OOM retries stay isolated across TaskSets and partial offers until completion") { + val clock = new ManualClock(1) + val scheduler = setupOomRetryScheduler(clock) + val manager = prepareOomRetries() + scheduler.submitTasks(FakeTask.createTaskSet(8, stageId = 0, stageAttemptId = 0)) + + val tasks = scheduler.resourceOffers(IndexedSeq( + WorkerOffer("a", "host-a", 4), WorkerOffer("b", "host-b", 4))).flatten + val retry = tasks.find(task => + scheduler.taskIdToTaskSetManager.get(task.taskId) eq manager).get + assert(retry.executorId === "a") + assert(retry.cpus === 1) + assert(tasks.count(_.executorId == "a") === 1) + assert(tasks.count(_.executorId == "b") === 4) + assert(scheduler.oomRetryReservationInfo === Some(oomReservationInfo(manager, "a"))) + + // The wait deadline must not end isolation after the retry has started. + clock.advance(60001) + assert(scheduler.resourceOffers(IndexedSeq(WorkerOffer("b", "host-b", 0)), + isAllFreeResources = false).flatten.isEmpty) + assert(scheduler.resourceOffers(IndexedSeq(WorkerOffer("a", "host-a", 3)), + isAllFreeResources = false).flatten.isEmpty) + + finishOomTestTask(retry) + assert(scheduler.oomRetryReservationInfo.isEmpty) + val ordinary = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("a", "host-a", 4))).flatten + assert(ordinary.size === 4) + assert(ordinary.forall(_.executorId == "a")) + } + + test("an OOM recovery launch clears the unschedulable TaskSet expiry") { + val scheduler = setupOomRetryScheduler(new ManualClock(1)) + val manager = prepareOomRetries() + scheduler.unschedulableTaskSetToExpiryTime(manager) = 60000L + + val tasks = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("a", "host-a", 4))).flatten + assert(tasks.size === 1) + assert(scheduler.taskIdToTaskSetManager.get(tasks.head.taskId) eq manager) + assert(scheduler.unschedulableTaskSetToExpiryTime.isEmpty) + } + + test("OOM isolation reselects an executor when the drained reservation cannot fit the task") { + val scheduler = setupOomRetryScheduler(new ManualClock(1)) + scheduler.submitTasks(FakeTask.createTaskSet(3)) + val blocker = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("small", "host-small", 1))).flatten.head + val otherBlockers = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("large", "host-large", 4))).flatten + assert(otherBlockers.size === 2) + + val profile = new TaskResourceProfile(new TaskResourceRequests().cpus(2).requests) + sc.resourceProfileManager.addResourceProfile(profile) + scheduler.submitTasks(FakeTask.createTaskSet(1, stageId = 1, stageAttemptId = 0, + priority = 0, rpId = profile.id)) + val manager = scheduler.taskSetManagerForAttempt(1, 0).get + (0 until 2).foreach { _ => + val tasks = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("origin", "host-origin", 2))).flatten + assert(tasks.size === 1) + failTaskWithOom(tasks.head) + } + scheduler.executorLost("origin", ExecutorProcessLost()) + + // The one-core executor is least busy, but cannot fit this two-core task after draining. + assert(scheduler.resourceOffers(IndexedSeq( + WorkerOffer("small", "host-small", 0), + WorkerOffer("large", "host-large", 2))).flatten.isEmpty) + finishOomTestTask(blocker) + otherBlockers.foreach(finishOomTestTask) + + // Reselect immediately, without advancing the isolation deadline. + val tasks = scheduler.resourceOffers(IndexedSeq( + WorkerOffer("small", "host-small", 1), + WorkerOffer("large", "host-large", 4))).flatten + assert(tasks.size === 1) + assert(scheduler.taskIdToTaskSetManager.get(tasks.head.taskId) eq manager) + assert(tasks.head.executorId === "large") + assert(tasks.head.cpus === 2) + } + + test("repeated OOM retries drain only one executor without killing existing tasks") { + val scheduler = setupOomRetryScheduler(new ManualClock(1)) + val backend = spy(new FakeSchedulerBackend) + scheduler.initialize(backend) + scheduler.submitTasks(FakeTask.createTaskSet(3, stageId = 0, stageAttemptId = 0)) + val blocker = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("a", "host-a", 1))).flatten.head + scheduler.resourceOffers(IndexedSeq(WorkerOffer("b", "host-b", 2))) + val manager = prepareOomRetries(numTasks = 2) + scheduler.submitTasks(FakeTask.createTaskSet(8, stageId = 2, stageAttemptId = 0)) + + // Select the least busy executor even though the busier executor is offered first. + val ordinary = scheduler.resourceOffers(IndexedSeq( + WorkerOffer("b", "host-b", 2), WorkerOffer("a", "host-a", 3))).flatten + assert(ordinary.size === 2) + assert(ordinary.forall(_.executorId == "b")) + assert(scheduler.oomRetryReservationInfo === Some(oomReservationInfo(manager, "a"))) + assert(scheduler.resourceOffers( + IndexedSeq(WorkerOffer("a", "host-a", 3))).flatten.isEmpty) + + finishOomTestTask(blocker) + // Keep the retry identity stable while the executor drains. + assert(scheduler.oomRetryReservationInfo === Some(oomReservationInfo(manager, "a"))) + val isolated = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("a", "host-a", 4))).flatten + assert(isolated.size === 1) + assert(scheduler.taskIdToTaskSetManager.get(isolated.head.taskId) eq manager) + assert(manager.copiesRunning.count(_ > 0) === 1) + val snapshotRead = new CountDownLatch(1) + val snapshotMatches = new AtomicBoolean(false) + val reader = new Thread(() => { + snapshotMatches.set( + scheduler.oomRetryReservationInfo.contains(oomReservationInfo(manager, "a"))) + snapshotRead.countDown() + }) + try { + scheduler.synchronized { + reader.start() + assert(snapshotRead.await(10, TimeUnit.SECONDS), "allocation must not take scheduler locks") + assert(snapshotMatches.get()) + } + } finally { + reader.join(10000) + } + // The second OOM retry must not drain another executor while this one is isolated. + val moreOrdinary = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("b", "host-b", 2))).flatten + assert(moreOrdinary.size === 2) + assert(moreOrdinary.forall(task => + scheduler.taskIdToTaskSetManager.get(task.taskId) ne manager)) + verify(backend, never()).killTask(anyLong(), anyString(), anyBoolean(), anyString()) + } + + test("OOM isolation waiting expires without rearming the reservation") { + val clock = new ManualClock(1) + val scheduler = setupOomRetryScheduler(clock, + config.SCHEDULER_OOM_RETRY_ISOLATION_TIMEOUT.key -> "1s") + scheduler.submitTasks(FakeTask.createTaskSet(1, stageId = 0, stageAttemptId = 0)) + scheduler.resourceOffers(IndexedSeq(WorkerOffer("a", "host-a", 1))) + val manager = prepareOomRetries() + scheduler.submitTasks(FakeTask.createTaskSet(2, stageId = 2, stageAttemptId = 0)) + + assert(scheduler.resourceOffers( + IndexedSeq(WorkerOffer("a", "host-a", 3))).flatten.isEmpty) + clock.advance(999) + assert(scheduler.resourceOffers( + IndexedSeq(WorkerOffer("a", "host-a", 3))).flatten.isEmpty) + clock.advance(1) + // Let the wait expire in an offer that cannot launch the retry yet. + assert(scheduler.resourceOffers( + IndexedSeq(WorkerOffer("a", "host-a", 0))).flatten.isEmpty) + assert(scheduler.oomRetryReservationInfo.isEmpty) + clock.advance(1) + val retry = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("a", "host-a", 1))).flatten + assert(retry.size === 1) + assert(scheduler.taskIdToTaskSetManager.get(retry.head.taskId) eq manager) + assert(retry.head.cpus === 1) + // Timed-out retries neither reacquire a reservation nor isolate their fallback attempt. + assert(scheduler.resourceOffers( + IndexedSeq(WorkerOffer("a", "host-a", 2))).flatten.size === 2) + } + + test("OOM isolation can expire while resource offers are being processed") { + val clock = new ManualClock(1) { + var tickOnRead = false + + override def getTimeMillis(): Long = { + val now = super.getTimeMillis() + if (tickOnRead) advance(1) + now + } + } + val scheduler = setupOomRetryScheduler(clock) + scheduler.submitTasks(FakeTask.createTaskSet(1, stageId = 0, stageAttemptId = 0)) + scheduler.resourceOffers(IndexedSeq(WorkerOffer("a", "host-a", 1))) + val manager = prepareOomRetries() + val offers = IndexedSeq(WorkerOffer("a", "host-a", 1)) + assert(scheduler.resourceOffers(offers).flatten.isEmpty) + + // Start three milliseconds before the deadline and cross it within one offer, rather + // than only advancing time between offers. The free core must not stay masked on expiry. + clock.advance(manager.oomRetryIsolationTimeRemaining(0) - 3) + clock.tickOnRead = true + val tasks = try { + scheduler.resourceOffers(offers).flatten + } finally { + clock.tickOnRead = false + } + assert(!manager.oomRetryNeedsIsolation(0)) + assert(tasks.size === 1) + assert(scheduler.taskIdToTaskSetManager.get(tasks.head.taskId) eq manager) + assert(tasks.head.executorId === "a") + } + + test("cancelling an isolated OOM retry protects its executor until the task exits") { + val scheduler = setupOomRetryScheduler(new ManualClock(1)) + val backend = mock[SchedulerBackend] + scheduler.initialize(backend) + val manager = prepareOomRetries() + val retry = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("a", "host-a", 4))).flatten.head + scheduler.submitTasks(FakeTask.createTaskSet(2, stageId = 0, stageAttemptId = 0)) + + scheduler.killAllTaskAttempts(1, interruptThread = true, reason = "test cancellation") + assert(scheduler.oomRetryReservationInfo === Some(oomReservationInfo(manager, "a"))) + verify(backend).killTask(retry.taskId, "a", true, "Stage cancelled: test cancellation") + assert(scheduler.resourceOffers( + IndexedSeq(WorkerOffer("a", "host-a", 3))).flatten.isEmpty) + + scheduler.synchronized { + failTask(retry.taskId, TaskState.KILLED, TaskKilled("test cancellation"), manager) + } + assert(scheduler.oomRetryReservationInfo.isEmpty) + assert(scheduler.resourceOffers( + IndexedSeq(WorkerOffer("a", "host-a", 4))).flatten.size === 2) + } + + test("pending OOM reservations clear when cancelled or completed by another attempt") { + val scheduler = setupOomRetryScheduler(new ManualClock(1)) + val backend = mock[SchedulerBackend] + scheduler.initialize(backend) + scheduler.submitTasks(FakeTask.createTaskSet(1, stageId = 0, stageAttemptId = 0)) + scheduler.resourceOffers(IndexedSeq(WorkerOffer("a", "host-a", 1))) + + Seq(false, true).zipWithIndex.foreach { case (partitionCompleted, i) => + val stageId = i + 1 + scheduler.submitTasks(FakeTask.createTaskSet(2, stageId, stageAttemptId = 0)) + val originals = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("origin", "origin-host", 2))).flatten + assert(originals.size === 2) + val otherPartition = originals(1) + failTaskWithOom(originals.head) + val retry = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("origin", "origin-host", 1))).flatten.head + failTaskWithOom(retry) + assert(scheduler.resourceOffers( + IndexedSeq(WorkerOffer("a", "host-a", 3))).flatten.isEmpty) + + scheduler.submitTasks(FakeTask.createTaskSet(1, stageId = i + 10, stageAttemptId = 0)) + assert(scheduler.resourceOffers( + IndexedSeq(WorkerOffer("a", "host-a", 1))).flatten.isEmpty) + scheduler.synchronized { + clearInvocations(backend) + if (partitionCompleted) { + scheduler.handlePartitionCompleted(stageId, originals.head.partitionId) + } else { + scheduler.killAllTaskAttempts( + stageId, interruptThread = true, reason = "test cancellation") + } + + // The other partition has not exited, so taskSetFinished has not released the reservation. + verify(backend).reviveOffers() + assert(scheduler.oomRetryReservationInfo.isEmpty) + } + val ordinary = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("a", "host-a", 1))).flatten + assert(ordinary.size === 1) + finishOomTestTask(ordinary.head) + finishOomTestTask(otherPartition) + } + } + + test("OOM reservations recover from decommissioning and executor loss") { + val scheduler = setupOomRetryScheduler(new ManualClock(1)) + scheduler.submitTasks(FakeTask.createTaskSet(2, stageId = 0, stageAttemptId = 0)) + scheduler.resourceOffers(IndexedSeq(WorkerOffer("a", "host-a", 1))) + val blocker = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("b", "host-b", 1))).flatten.head + val manager = prepareOomRetries() + + assert(scheduler.resourceOffers( + IndexedSeq(WorkerOffer("a", "host-a", 3))).flatten.isEmpty) + scheduler.executorDecommission("a", ExecutorDecommissionInfo("test", None)) + assert(scheduler.resourceOffers( + IndexedSeq(WorkerOffer("b", "host-b", 3))).flatten.isEmpty) + finishOomTestTask(blocker) + val retry = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("b", "host-b", 4))).flatten + assert(retry.size === 1) + assert(scheduler.taskIdToTaskSetManager.get(retry.head.taskId) eq manager) + + scheduler.executorLost("b", ExecutorExited(0, exitCausedByApp = false, "test loss")) + assert(scheduler.oomRetryReservationInfo.isEmpty) + val replacement = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("c", "host-c", 4))).flatten + assert(replacement.size === 1) + assert(scheduler.taskIdToTaskSetManager.get(replacement.head.taskId) eq manager) + assert(replacement.head.executorId === "c") + } + + test("OOM retries preserve exclusions, resource profiles and custom resource requirements") { + val scheduler = setupSchedulerWithMockTaskSetExcludelist( + config.SCHEDULER_OOM_RETRY_ENABLED.key -> "true") + val execReqs = new ExecutorResourceRequests().cores(4).resource(GPU, 1) + val taskReqs = new TaskResourceRequests().cpus(1).resource(GPU, 0.5) + val rp = new ResourceProfile(execReqs.requests, taskReqs.requests) + scheduler.sc.resourceProfileManager.addResourceProfile(rp) + def offer(execId: String, gpu: Double, profileId: Int = rp.id): WorkerOffer = { + WorkerOffer(execId, s"host-$execId", 4, resources = new ExecutorResourcesAmounts( + Map(GPU -> toInternalResource(Map("0" -> gpu)))), resourceProfileId = profileId) + } + scheduler.submitTasks(FakeTask.createTaskSet(1, stageId = 1, stageAttemptId = 0, + priority = 0, rpId = rp.id)) + (0 until 2).foreach { _ => + val tasks = scheduler.resourceOffers(IndexedSeq(offer("origin", 1))).flatten + assert(tasks.size === 1) + failTaskWithOom(tasks.head) + } + val excluded = stageToMockTaskSetExcludelist(1) + when(excluded.isNodeExcludedForTaskSet("host-node")).thenReturn(true) + when(excluded.isExecutorExcludedForTaskSet("executor")).thenReturn(true) + when(excluded.isExecutorExcludedForTask("task", 0)).thenReturn(true) + + val tasks = scheduler.resourceOffers(IndexedSeq( + offer("wrong-profile", 1, ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID), + offer("no-gpu", 0), offer("fractional", 0.25), + offer("node", 1), offer("executor", 1), offer("task", 1), + offer("eligible", 0.5))).flatten + assert(tasks.map(_.executorId) === Seq("eligible")) + assert(tasks.head.cpus === 1) + assert(tasks.head.resources(GPU).values.sum === ONE_ENTIRE_RESOURCE / 2) + } + + test("OOM retry placement bypasses locality without changing ordinary locality") { + val scheduler = setupOomRetryScheduler(new ManualClock(1), + config.LOCALITY_WAIT.key -> "60s") + scheduler.resourceOffers(IndexedSeq( + WorkerOffer("a", "host-a", 0), WorkerOffer("b", "host-b", 0))) + scheduler.submitTasks(FakeTask.createTaskSet(1, stageId = 1, stageAttemptId = 0, + Seq(TaskLocation("host-a", "a")))) + val manager = scheduler.taskSetManagerForAttempt(1, 0).get + val original = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("a", "host-a", 1))).flatten.head + failTaskWithOom(original) + scheduler.submitTasks(FakeTask.createTaskSet(1, stageId = 0, stageAttemptId = 0, + Seq(TaskLocation("host-a", "a")))) + + val retry = scheduler.resourceOffers(IndexedSeq( + WorkerOffer("a", "host-a", 0), WorkerOffer("b", "host-b", 2))).flatten + assert(retry.size === 1) + assert(retry.head.executorId === "b") + assert(manager.taskInfos(retry.head.taskId).taskLocality === TaskLocality.ANY) + val ordinary = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("a", "host-a", 1))).flatten + assert(ordinary.size === 1) + val ordinaryManager = scheduler.taskSetManagerForAttempt(0, 0).get + assert(ordinaryManager.taskInfos(ordinary.head.taskId).taskLocality === + TaskLocality.PROCESS_LOCAL) + } + + test("OOM-affected tasks are not speculated but ordinary tasks still are") { + val clock = new ManualClock(1) + val scheduler = setupOomRetryScheduler(clock, + config.SPECULATION_ENABLED.key -> "true", + config.SPECULATION_TASK_DURATION_THRESHOLD.key -> "1ms", + config.EXECUTOR_CORES.key -> "8") + scheduler.submitTasks(FakeTask.createTaskSet(2)) + val manager = scheduler.taskSetManagerForAttempt(0, 0).get + val originals = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("a", "host-a", 2))).flatten + // Leave an entry in the speculative pending queues from before the OOM. + manager.speculatableTasks += originals.head.index + manager.addPendingTask(originals.head.index, speculatable = true) + failTaskWithOom(originals.head) + val retry = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("b", "host-b", 1))).flatten.head + + clock.advance(1000) + assert(manager.checkSpeculatableTasks(0)) + assert(manager.speculatableTasks.contains(originals(1).index)) + val speculative = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("c", "host-c", 2))).flatten + assert(speculative.map(_.index) === Seq(originals(1).index)) + assert(manager.taskInfos(speculative.head.taskId).speculative) + assert(manager.copiesRunning(retry.index) === 1) + } + + test("OOM retries leave normal scheduling unchanged when recovery is disabled by default") { + val scheduler = setupScheduler() + val manager = prepareOomRetries() + scheduler.submitTasks(FakeTask.createTaskSet(3, stageId = 0, stageAttemptId = 0)) + val tasks = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("a", "host-a", 4))).flatten + assert(tasks.size === 4) + assert(tasks.exists(task => scheduler.taskIdToTaskSetManager.get(task.taskId) eq manager)) + assert(tasks.forall(_.cpus == 1)) + } + + test("barrier scheduling respects OOM isolation and does not retry individual OOM failures") { + val scheduler = setupOomRetryScheduler(new ManualClock(1)) + val manager = prepareOomRetries() + scheduler.submitTasks(FakeTask.createBarrierTaskSet(2)) + val tasks = scheduler.resourceOffers(IndexedSeq( + WorkerOffer("a", "host-a", 4, Some("host-a:1000")), + WorkerOffer("b", "host-b", 1, Some("host-b:1000")))).flatten + assert(tasks.size === 1) + assert(scheduler.taskIdToTaskSetManager.get(tasks.head.taskId) eq manager) + + val barrier = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("b", "host-b", 2, Some("host-b:1000")))).flatten + assert(barrier.size === 2) + assert(barrier.forall(_.executorId == "b")) + val barrierManager = scheduler.taskSetManagerForAttempt(0, 0).get + failTaskWithOom(barrier.head) + assert(barrierManager.isZombie) + + scheduler.submitTasks(FakeTask.createTaskSet(1, stageId = 2, stageAttemptId = 0)) + val ordinary = scheduler.resourceOffers( + IndexedSeq(WorkerOffer("b", "host-b", 1))).flatten + assert(ordinary.size === 1) + assert(scheduler.taskIdToTaskSetManager.get(ordinary.head.taskId) ne barrierManager) + } + + for (oomRecovery <- Seq(false, true); failedPartition <- 0 until 3) { + test(s"barrier serialization failure preserves other launches: " + + s"OOM recovery=$oomRecovery, partition=$failedPartition") { + val clock = new ManualClock(1) + val scheduler = setupOomRetryScheduler(clock, + config.SCHEDULER_OOM_RETRY_ENABLED.key -> oomRecovery.toString, + config.EXECUTOR_CORES.key -> "8", + config.CPUS_PER_TASK.key -> "2", + EXECUTOR_GPU_ID.amountConf -> "2", + TASK_GPU_ID.amountConf -> "0.5") + def offer(executorId: String, cores: Int): WorkerOffer = { + val resources = new ExecutorResourcesAmounts( + Map(GPU -> Map("0" -> ONE_ENTIRE_RESOURCE, "1" -> ONE_ENTIRE_RESOURCE))) + WorkerOffer(executorId, s"host-$executorId", cores, + Some(s"host-$executorId:1000"), resources) + } + scheduler.submitTasks(FakeTask.createTaskSet(1, stageId = 1, stageAttemptId = 0, + priority = 1, rpId = ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)) + val retryManager = scheduler.taskSetManagerForAttempt(1, 0).get + (0 until 2).foreach { _ => + val original = scheduler.resourceOffers(IndexedSeq(offer("origin", 2))).flatten.head + failTaskWithOom(original) + } + + // The RDD/function broadcast serializes, but the real ResultTask's partition does not. + // Try each bad partition so both early and later preparation failures are covered without + // depending on the iteration order of barrierPendingLaunchTasks. + val data = (0 until 3).map[AnyRef] { i => + if (i == failedPartition) new Object else Int.box(i) + } + val rdd = sc.parallelize(data, 3).barrier().mapPartitions(_.map(_ => 1)) + val func = (_: TaskContext, values: Iterator[Int]) => values.sum + val serializer = sc.env.closureSerializer.newInstance() + val serialized = serializer.serialize((rdd, func): AnyRef) + val binary = new Array[Byte](serialized.remaining()) + serialized.get(binary) + val broadcast = sc.broadcast(binary) + val metrics = TaskMetrics.registered + val serializedMetrics = serializer.serialize(metrics).array() + val barrierTasks = rdd.partitions.map[Task[_]] { partition => + new ResultTask[Int, Int](0, 0, broadcast, partition, 3, Nil, partition.index, + JobArtifactSet.defaultJobArtifactSet, new Properties, serializedMetrics, isBarrier = true) + } + scheduler.submitTasks(new TaskSet(barrierTasks, 0, 0, 0, null, + ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID, None)) + val barrierManager = scheduler.taskSetManagerForAttempt(0, 0).get + scheduler.submitTasks(FakeTask.createTaskSet(5, stageId = 2, stageAttemptId = 0, + priority = 2, rpId = ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)) + assert(scheduler.rootPool.getSortedTaskSetQueue.head eq barrierManager) + + val offers = IndexedSeq(offer("a", 4), offer("b", 6)) + val tasks = scheduler.resourceOffers(offers).flatten + assert(failedTaskSet) + assert(failedTaskSetReason.contains("Failed to serialize task")) + assert(barrierManager.isZombie) + assert(barrierManager.runningTasks === 0) + assert(barrierManager.barrierPendingLaunchTasks.isEmpty) + assert(scheduler.taskSetManagerForAttempt(0, 0).isEmpty) + val prepared = barrierManager.taskInfos.values.filter(_.index != failedPartition) + assert(prepared.forall(_.killed)) + assert(!scheduler.taskIdToTaskSetManager.containsValue(barrierManager)) + + val retries = tasks.filter(task => + scheduler.taskIdToTaskSetManager.get(task.taskId) eq retryManager) + assert(retries.size === 1) + assert(retries.head.attemptNumber === 2) + assert(tasks.size === (if (oomRecovery) 4 else 5)) + assert(tasks.forall(_.cpus == 2)) + assert(tasks.count(_.executorId == "b") === 3) + assert(scheduler.rootPool.runningTasks === tasks.size) + offers.foreach { offered => + val used = tasks.filter(_.executorId == offered.executorId) + .map(_.resources(GPU).values.sum).sum + assert(offered.resources.availableResources(GPU).values.sum === + 2.0 - ResourceAmountUtils.toFractionalResource(used)) + } + + // The returned retry remains isolated until its real completion, then releases capacity. + clock.advance(60001) + if (oomRecovery) { + assert(scheduler.resourceOffers(IndexedSeq(offer("a", 2))).flatten.isEmpty) + } + finishOomTestTask(retries.head) + val remaining = scheduler.resourceOffers(IndexedSeq(offer("a", 4))).flatten + assert(remaining.size === (if (oomRecovery) 2 else 1)) + } + } + + test("legacy barrier abort preserves other launches and releases its assignments") { + val scheduler = setupSchedulerWithMockTaskSetExcludelist( + config.LEGACY_LOCALITY_WAIT_RESET.key -> "true") + scheduler.submitTasks(FakeTask.createTaskSet(1, stageId = 0, stageAttemptId = 0)) + scheduler.submitTasks(FakeTask.createBarrierTaskSet(2, 1, 0, 0, + ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)) + scheduler.submitTasks(FakeTask.createTaskSet(2, stageId = 2, stageAttemptId = 0)) + when(stageToMockTaskSetExcludelist(0).isExecutorExcludedForTaskSet("b")).thenReturn(true) + when(stageToMockTaskSetExcludelist(1).isExecutorExcludedForTaskSet("b")).thenReturn(true) + val barrierManager = scheduler.taskSetManagerForAttempt(1, 0).get + + // After stage 0 launches, two slots remain, but the barrier can only use one of them. + val tasks = scheduler.resourceOffers(IndexedSeq( + WorkerOffer("a", "host-a", 2, Some("host-a:1000")), + WorkerOffer("b", "host-b", 1, Some("host-b:1000")))).flatten + assert(failedTaskSet) + assert(failedTaskSetReason.contains("of 2 tasks got resource offers")) + assert(tasks.map(task => scheduler.taskIdToTaskSetManager.get(task.taskId).stageId).sorted === + Seq(0, 2, 2)) + assert(tasks.map(_.executorId).sorted === Seq("a", "a", "b")) + assert(barrierManager.runningTasks === 0) + assert(barrierManager.barrierPendingLaunchTasks.isEmpty) + assert(scheduler.taskSetManagerForAttempt(1, 0).isEmpty) + } + + Seq(1.0, 0.5).foreach { gpuPerTask => + test(s"barrier scheduling accounts for GPU $gpuPerTask consumed by a first OOM retry") { + val scheduler = setupOomRetryScheduler(new ManualClock(1), + config.EXECUTOR_CORES.key -> "8", + EXECUTOR_GPU_ID.amountConf -> "2", + TASK_GPU_ID.amountConf -> gpuPerTask.toString, + config.LEGACY_LOCALITY_WAIT_RESET.key -> "true") + def offer(): WorkerOffer = { + val resources = new ExecutorResourcesAmounts( + Map(GPU -> Map("0" -> ONE_ENTIRE_RESOURCE, "1" -> ONE_ENTIRE_RESOURCE))) + WorkerOffer("gpu", "host-gpu", 8, Some("host-gpu:1000"), resources) + } + scheduler.submitTasks(FakeTask.createTaskSet(1, stageId = 1, stageAttemptId = 0)) + val original = scheduler.resourceOffers(IndexedSeq(offer())).flatten.head + failTaskWithOom(original) + + val numBarrierTasks = (2 / gpuPerTask).toInt + scheduler.submitTasks(FakeTask.createBarrierTaskSet(numBarrierTasks)) + val barrierManager = scheduler.taskSetManagerForAttempt(0, 0).get + val tasks = scheduler.resourceOffers(IndexedSeq(offer())).flatten + assert(tasks.size === 1) + assert(tasks.head.attemptNumber === 1) + assert(scheduler.taskIdToTaskSetManager.get(tasks.head.taskId) ne barrierManager) + assert(!barrierManager.isZombie) + assert(!failedTaskSet) + + finishOomTestTask(tasks.head) + val barrierTasks = scheduler.resourceOffers(IndexedSeq(offer())).flatten + assert(barrierTasks.size === numBarrierTasks) + assert(barrierTasks.forall(task => + scheduler.taskIdToTaskSetManager.get(task.taskId) eq barrierManager)) + } + } + + test("GPU-limited barrier stages do not count resources on an OOM-reserved executor") { + val scheduler = setupOomRetryScheduler(new ManualClock(1), + config.EXECUTOR_CORES.key -> "4", + EXECUTOR_GPU_ID.amountConf -> "2", + TASK_GPU_ID.amountConf -> "1", + config.LEGACY_LOCALITY_WAIT_RESET.key -> "true") + def offer(execId: String, cores: Int, addresses: Seq[String]): WorkerOffer = { + val resources = new ExecutorResourcesAmounts( + Map(GPU -> addresses.map(_ -> ONE_ENTIRE_RESOURCE).toMap)) + WorkerOffer(execId, s"host-$execId", cores, Some(s"host-$execId:1000"), resources) + } + scheduler.submitTasks(FakeTask.createTaskSet(1, stageId = 1, stageAttemptId = 0)) + (0 until 2).foreach { _ => + val tasks = scheduler.resourceOffers( + IndexedSeq(offer("origin", 4, Seq("0", "1")))).flatten + assert(tasks.size === 1) + failTaskWithOom(tasks.head) + } + assert(scheduler.resourceOffers( + IndexedSeq(offer("a", 4, Seq("0", "1")))).flatten.size === 1) + scheduler.submitTasks(FakeTask.createBarrierTaskSet(3)) + val barrierManager = scheduler.taskSetManagerForAttempt(0, 0).get + + // The reserved executor has one free GPU, but it must not count toward barrier capacity. + // Otherwise legacy delay scheduling aborts after partially assigning the other two tasks. + val tasks = scheduler.resourceOffers(IndexedSeq( + offer("a", 3, Seq("1")), offer("b", 4, Seq("0", "1")))).flatten + assert(tasks.isEmpty) + assert(!barrierManager.isZombie) + assert(!failedTaskSet) + + val barrier = scheduler.resourceOffers(IndexedSeq( + offer("a", 3, Seq("1")), offer("b", 4, Seq("0", "1")), + offer("c", 4, Seq("0", "1")))).flatten + assert(barrier.size === 3) + assert(barrier.forall(_.executorId != "a")) + } + test("Scheduler correctly accounts for multiple CPUs per task") { val taskCpus = 2 val taskScheduler = setupSchedulerWithMaster( diff --git a/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala index 1e3ae50ce6cdf..8935aba4851cd 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala @@ -45,7 +45,7 @@ import org.apache.spark.resource.TestResourceIDs._ import org.apache.spark.scheduler.cluster.CoarseGrainedSchedulerBackend import org.apache.spark.serializer.SerializerInstance import org.apache.spark.storage.BlockManagerId -import org.apache.spark.util.{AccumulatorV2, Clock, ManualClock, SystemClock} +import org.apache.spark.util.{AccumulatorV2, Clock, ManualClock, SparkExitCode, SystemClock} import org.apache.spark.util.ArrayImplicits._ class FakeDAGScheduler(sc: SparkContext, taskScheduler: FakeTaskScheduler) @@ -705,6 +705,85 @@ class TaskSetManagerSuite assert(sched.taskSetsFailed.contains(taskSet.id)) } + test("OOM retries use structured executor loss reasons") { + val conf = new SparkConf().set(config.SCHEDULER_OOM_RETRY_ENABLED, true) + sc = new SparkContext("local", "test", conf) + val clock = new ManualClock(1) + sched = new FakeTaskScheduler(sc, clock) + val containerOom = ExecutorExited(137, exitCausedByApp = true) + containerOom.isOutOfMemoryError = true + val reasons = Seq( + (ExecutorExited(SparkExitCode.OOM, exitCausedByApp = true), true), + (containerOom, true), + (ExecutorExited(137, exitCausedByApp = true), false), + (ExecutorExited(SparkExitCode.OOM, exitCausedByApp = false), false)) + + for (((reason, expectedOom), stageId) <- reasons.zipWithIndex) { + withClue(s"executor loss $reason: ") { + val execId = s"exec$stageId" + sched.addExecutor(execId, "host1") + val manager = new TaskSetManager(sched, + FakeTask.createTaskSet(1, stageId, stageAttemptId = 0), + MAX_TASK_FAILURES, clock = clock) + val task = manager.resourceOffer(execId, "host1", ANY)._1.get + manager.taskInfos(task.taskId).launchSucceeded() + sched.removeExecutor(execId) + manager.executorLost(execId, "host1", reason) + + val failure = sched.endedTasks(task.index).asInstanceOf[ExecutorLostFailure] + assert(failure.isOutOfMemoryError === expectedOom) + assert(manager.isPendingOomRetry(task.index) === expectedOom) + assert(!manager.oomRetryNeedsIsolation(task.index)) + } + } + } + + for (exceptionFirst <- Seq(true, false)) { + test(s"OOM retries count each failed attempt once (exception first: $exceptionFirst)") { + val conf = new SparkConf().set(config.SCHEDULER_OOM_RETRY_ENABLED, true) + sc = new SparkContext("local", "test", conf) + val clock = new ManualClock(1) + sched = new FakeTaskScheduler(sc, clock, ("exec1", "host1"), ("exec2", "host2")) + val manager = new TaskSetManager( + sched, FakeTask.createTaskSet(2), MAX_TASK_FAILURES, clock = clock) + val first = manager.resourceOffer("exec1", "host1", ANY)._1.get + manager.taskInfos(first.taskId).launchSucceeded() + val failure = new ExceptionFailure( + new RuntimeException("spill failed", new OutOfMemoryError("heap")), Nil) + + def loseExecutor(): Unit = { + sched.removeExecutor("exec1") + manager.executorLost("exec1", "host1", + ExecutorExited(SparkExitCode.OOM, exitCausedByApp = true)) + } + + if (exceptionFirst) { + manager.handleFailedTask(first.taskId, TaskState.FAILED, failure) + loseExecutor() + } else { + loseExecutor() + manager.handleFailedTask(first.taskId, TaskState.FAILED, failure) + } + assert(manager.isPendingOomRetry(first.index)) + assert(!manager.oomRetryNeedsIsolation(first.index)) + assert(manager.copiesRunning(first.index) === 0) + + val second = manager.resourceOffer("exec2", "host2", ANY)._1.get + assert(second.index === first.index) + manager.taskInfos(second.taskId).launchSucceeded() + manager.handleFailedTask(second.taskId, TaskState.FAILED, failure) + assert(manager.isPendingOomRetry(first.index)) + assert(manager.oomRetryNeedsIsolation(first.index)) + + val third = manager.resourceOfferOomRetry( + first.index, "exec2", "host2", taskCpus = 1, taskResources = Map.empty).get + manager.handleSuccessfulTask(third.taskId, createTaskResult(0)) + assert(manager.successful(first.index)) + assert(!manager.isPendingOomRetry(first.index)) + assert(!manager.oomRetryNeedsIsolation(first.index)) + } + } + test("SPARK-31837: Shift to the new highest locality level if there is when recomputeLocality") { sc = new SparkContext("local", "test") sched = new FakeTaskScheduler(sc) @@ -3074,6 +3153,21 @@ class TaskSetManagerSuite "a pipelined task set must abort after the first task failure") } + test("OOM recovery does not retry individual tasks in a pipelined task set") { + sc = new SparkContext(new SparkConf().setMaster("local").setAppName("test") + .set(config.SCHEDULER_OOM_RETRY_ENABLED, true)) + sched = new FakeTaskScheduler(sc, ("exec1", "host1")) + val taskSet = pipelinedTaskSet() + val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES) + val task = manager.resourceOffer("exec1", "host1", ANY)._1.get + + manager.handleFailedTask(task.taskId, TaskState.FAILED, + new ExceptionFailure(new OutOfMemoryError("test OOM"), Nil)) + + assert(sched.taskSetsFailed.contains(taskSet.id)) + assert(manager.pendingOomRetries.isEmpty) + } + test("pipelined task set counts an executor-loss failure that would normally be uncounted") { sc = new SparkContext("local", "test") sched = new FakeTaskScheduler(sc, ("exec1", "host1")) diff --git a/core/src/test/scala/org/apache/spark/util/JsonProtocolSuite.scala b/core/src/test/scala/org/apache/spark/util/JsonProtocolSuite.scala index 4639aa4da9dc2..f0b1fc20f8fe5 100644 --- a/core/src/test/scala/org/apache/spark/util/JsonProtocolSuite.scala +++ b/core/src/test/scala/org/apache/spark/util/JsonProtocolSuite.scala @@ -305,6 +305,37 @@ class JsonProtocolSuite extends SparkFunSuite { assertEquals(exceptionFailure, JsonProtocol.taskEndReasonFromJson(oldEvent)) } + test("OOM task failure classification is preserved in event logs") { + val exceptionFailure = new ExceptionFailure( + new RuntimeException("spill failed", new OutOfMemoryError("heap")), Nil) + testTaskEndReason(exceptionFailure) + val executorFailure = ExecutorLostFailure("1", true, Some("exit code 137")) + executorFailure.isOutOfMemoryError = true + testTaskEndReason(executorFailure) + } + + test("OOM task failure classification is compatible with old event logs") { + val wrapped = new ExceptionFailure( + new RuntimeException("spill failed", new OutOfMemoryError("heap")), Nil) + val wrappedJson = toJsonString(JsonProtocol.taskEndReasonToJson(wrapped, _)) + .removeField("Out Of Memory Error") + assert(!JsonProtocol.taskEndReasonFromJson(wrappedJson) + .asInstanceOf[ExceptionFailure].isOutOfMemoryError) + + val direct = new ExceptionFailure(new OutOfMemoryError("heap"), Nil) + val directJson = toJsonString(JsonProtocol.taskEndReasonToJson(direct, _)) + .removeField("Out Of Memory Error") + assert(JsonProtocol.taskEndReasonFromJson(directJson) + .asInstanceOf[ExceptionFailure].isOutOfMemoryError) + + val executorFailure = ExecutorLostFailure("1", true, Some("OOMKilled, exit code 137")) + executorFailure.isOutOfMemoryError = true + val executorJson = toJsonString(JsonProtocol.taskEndReasonToJson(executorFailure, _)) + .removeField("Out Of Memory Error") + assert(!JsonProtocol.taskEndReasonFromJson(executorJson) + .asInstanceOf[ExecutorLostFailure].isOutOfMemoryError) + } + test("StageInfo backward compatibility (details, accumulables)") { val info = makeStageInfo(1, 2, 3, 4L, 5L) val newJson = toJsonString( @@ -1392,6 +1423,7 @@ private[spark] object JsonProtocolSuite extends Assertions { assert(r1.description === r2.description) assertSeqEquals(r1.stackTrace, r2.stackTrace, assertStackTraceElementEquals) assert(r1.fullStackTrace === r2.fullStackTrace) + assert(r1.isOutOfMemoryError === r2.isOutOfMemoryError) val filteredUpdates = r1.accumUpdates .filterNot { acc => acc.name.exists(accumulableExcludeList.contains) } assertSeqEquals[AccumulableInfo](filteredUpdates, r2.accumUpdates, (a, b) => a.equals(b)) @@ -1403,11 +1435,11 @@ private[spark] object JsonProtocolSuite extends Assertions { assert(jobId1 === jobId2) assert(partitionId1 === partitionId2) assert(attemptNumber1 === attemptNumber2) - case (ExecutorLostFailure(execId1, exit1CausedByApp, reason1), - ExecutorLostFailure(execId2, exit2CausedByApp, reason2)) => - assert(execId1 === execId2) - assert(exit1CausedByApp === exit2CausedByApp) - assert(reason1 === reason2) + case (r1: ExecutorLostFailure, r2: ExecutorLostFailure) => + assert(r1.execId === r2.execId) + assert(r1.exitCausedByApp === r2.exitCausedByApp) + assert(r1.reason === r2.reason) + assert(r1.isOutOfMemoryError === r2.isOutOfMemoryError) case (ExecutorShutdownFailure(execId1), ExecutorShutdownFailure(execId2)) => assert(execId1 === execId2) case (UnknownReason, UnknownReason) => diff --git a/docs/configuration.md b/docs/configuration.md index 683d1048dca60..21ad11e9faa04 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -2949,6 +2949,42 @@ Apart from these, the following properties are also available, and may be useful 0.8.0 + + spark.scheduler.oomRetry.enabled + false + + Prefer eligible idle executors when retrying a task after an out-of-memory failure. + After two OOM failures of the same task, temporarily reserve one executor for its retry: + existing tasks finish first, and no other tasks are placed there until the retry exits. + At most one executor per application is reserved. Other repeated OOM retries wait up to + spark.scheduler.oomRetry.isolationTimeout before falling back to ordinary placement. + Normal tasks retain their ordinary placement policy. OOM retries may ignore preferred + locations, but still respect exclusions, resource profiles, and custom resource requirements. + Task CPU requests, executor memory limits, and spark.task.maxFailures are unchanged. + Dynamic allocation counts the reserved executor separately from the remaining work, within + the configured executor limits, so its unused slots do not suppress requests for other tasks. + Barrier and pipelined stages are excluded, and OOM-affected tasks are not speculated. + Recognized failures include JVM/Spark OOM exceptions (including wrapped exceptions), Spark + executor exit code 52, and Kubernetes executor-container OOMKilled termination. + Exit code 137 alone is not considered OOM. Native errors without a typed OOM signal + are not recognized; fixed per-task native memory limits are not increased. + An idle executor may retain cached or native memory, so recovery is not guaranteed. + + 5.0.0 + + + spark.scheduler.oomRetry.isolationTimeout + 60s + + Maximum wait for isolation, measured from the task's latest OOM failure, when + spark.scheduler.oomRetry.enabled is enabled. If an executor cannot be reserved + and drained within this time, the pending retry falls back to ordinary placement. The same + pending retry does not start another wait; a new OOM failure starts a new wait. + This timeout does not end isolation or kill a retry that has already started. + Must be positive. + + 5.0.0 + spark.scheduler.revive.interval 1s diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorPodsLifecycleManager.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorPodsLifecycleManager.scala index b914c3054d209..82e5590b211ba 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorPodsLifecycleManager.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorPodsLifecycleManager.scala @@ -22,7 +22,7 @@ import scala.collection.mutable import scala.jdk.CollectionConverters._ import com.google.common.cache.CacheBuilder -import io.fabric8.kubernetes.api.model.{Pod, PodBuilder} +import io.fabric8.kubernetes.api.model.{ContainerStateTerminated, Pod, PodBuilder} import io.fabric8.kubernetes.client.KubernetesClient import io.fabric8.kubernetes.client.dsl.base.{PatchContext, PatchType} @@ -269,7 +269,8 @@ private[spark] class ExecutorPodsLifecycleManager( } private def findExitReason(podState: FinalPodState, execId: Long): ExecutorExited = { - val exitCode = findExitCode(podState) + val executorTermination = findExecutorTermination(podState) + val exitCode = executorTermination.map(_.getExitCode.toInt).getOrElse(UNKNOWN_EXIT_CODE) val (exitCausedByApp, exitMessage) = podState match { case PodDeleted(_) => (false, s"The executor with id $execId was deleted by a user or the framework.") @@ -277,7 +278,11 @@ private[spark] class ExecutorPodsLifecycleManager( val msg = exitReasonMessage(podState, execId, exitCode) (true, msg) } - ExecutorExited(exitCode, exitCausedByApp, exitMessage) + val exitReason = ExecutorExited(exitCode, exitCausedByApp, exitMessage) + // Exit code 137 alone can also mean a non-OOM SIGKILL. Only trust the executor container. + exitReason.isOutOfMemoryError = exitCausedByApp && + (exitReason.isOutOfMemoryError || executorTermination.exists(_.getReason == "OOMKilled")) + exitReason } private def exitReasonMessage(podState: FinalPodState, execId: Long, exitCode: Int) = { @@ -301,13 +306,12 @@ private[spark] class ExecutorPodsLifecycleManager( """.stripMargin } - private def findExitCode(podState: FinalPodState): Int = { + private def findExecutorTermination(podState: FinalPodState): Option[ContainerStateTerminated] = { podState.pod.getStatus.getContainerStatuses.asScala.find { containerStatus => - containerStatus.getName == sparkContainerName && - containerStatus.getState.getTerminated != null - }.map { terminatedContainer => - terminatedContainer.getState.getTerminated.getExitCode.toInt - }.getOrElse(UNKNOWN_EXIT_CODE) + containerStatus.getName == sparkContainerName + }.flatMap { containerStatus => + Option(containerStatus.getState).flatMap(state => Option(state.getTerminated)) + } } private def isPodInactive(pod: Pod): Boolean = { diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorPodsLifecycleManagerSuite.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorPodsLifecycleManagerSuite.scala index 00f7fb227bbd0..308334e764dc9 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorPodsLifecycleManagerSuite.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorPodsLifecycleManagerSuite.scala @@ -23,8 +23,7 @@ import io.fabric8.kubernetes.client.KubernetesClient import io.fabric8.kubernetes.client.dsl.PodResource import io.fabric8.kubernetes.client.dsl.base.PatchContext import org.mockito.{ArgumentCaptor, Mock, MockitoAnnotations} -import org.mockito.ArgumentMatchers.any -import org.mockito.ArgumentMatchers.anyString +import org.mockito.ArgumentMatchers.{any, anyString, eq => meq} import org.mockito.Mockito.{mock, never, times, verify, when} import org.mockito.invocation.InvocationOnMock import org.mockito.stubbing.Answer @@ -231,6 +230,35 @@ class ExecutorPodsLifecycleManagerSuite extends SparkFunSuite with BeforeAndAfte verify(schedulerBackend).doRemoveExecutor("1", expectedLossReason) } + for ((name, exitCode, executorReason, sidecarReason, deleted, expectedOom) <- Seq( + ("executor OOMKilled", 137, "OOMKilled", "Error", false, true), + ("JVM OOM exit", SparkExitCode.OOM, "Error", "Error", false, true), + ("unclassified SIGKILL", 137, "Error", "Error", false, false), + ("sidecar OOMKilled", 137, "Error", "OOMKilled", false, false), + ("deleted executor OOMKilled", 137, "OOMKilled", "Error", true, false), + ("deleted JVM OOM exit", SparkExitCode.OOM, "Error", "Error", true, false))) { + test(s"OOM classification uses the executor termination: $name") { + // This fixture uses a custom executor container name and lists the sidecar first. + val failedPod = failedExecutorWithSidecarStatusListedFirst(1) + val statuses = failedPod.getStatus.getContainerStatuses + statuses.get(0).getState.getTerminated.setReason(sidecarReason) + statuses.get(1).getState.getTerminated.setExitCode(exitCode) + statuses.get(1).getState.getTerminated.setReason(executorReason) + // A pod-level reason must not override the executor container's actual termination. + failedPod.getStatus.setReason("OOMKilled") + if (deleted) { + failedPod.getMetadata.setDeletionTimestamp("2026-08-25T00:00:00Z") + } + snapshotsStore.updatePod(failedPod) + snapshotsStore.notifySubscribers() + val lossReason = ArgumentCaptor.forClass(classOf[ExecutorExited]) + verify(schedulerBackend).doRemoveExecutor(meq("1"), lossReason.capture()) + assert(lossReason.getValue.exitCode == exitCode) + assert(lossReason.getValue.exitCausedByApp == !deleted) + assert(lossReason.getValue.isOutOfMemoryError == expectedOom) + } + } + test("Don't delete pod from K8s if deletionTimestamp is already set.") { // Create a failed pod with deletionTimestamp already in the past val basePod = failedExecutorWithoutDeletion(1)