[SPARK-58187][CORE] Increase task CPUs on OOM retry to reduce concurrency and grow per-task memory

[ulysses-you]

What changes were proposed in this pull request?

When a task fails with out-of-memory, the retry can now request more CPUs than the ResourceProfile's spark.task.cpus, controlled by a new config spark.task.oomRetryCpusIncrement (default 0 = disabled). Each OOM retry of a task gets spark.task.cpus + spark.task.oomRetryCpusIncrement * N CPUs (N = the number of OOM failures of that task), capped at the executor's total core count.

Two OOM sources are recognized:

• SparkOutOfMemoryError (execution memory pool exhaustion), via the ExceptionFailure path in handleFailedTask.
• A JVM heap OutOfMemoryError that kills the executor, recognized by the executor exit code in executorLost (configurable via spark.task.oomRetryExecutorExitCodes, default 52).

The extra CPUs reduce the number of concurrent tasks on the executor, which increases the retrying task's share of the execution memory pool (the pool is split by the number of active tasks). If the offered executor does not have enough free CPUs, the retry waits for a sufficient offer rather than falling back. Barrier stages are excluded. Speculative copies reuse the same per-index retry count, so they request the same CPUs as the running OOM retry.

Why are the changes needed?

Today an OOM-failed task is retried with identical resources, which is unlikely to succeed and wastes retry attempts. Giving the retry more CPUs (hence more memory) improves the chance of success.

Does this PR introduce any user-facing change?

Yes, two new configs spark.task.oomRetryCpusIncrement and spark.task.oomRetryExecutorExitCodes, both documented in configuration.md. Disabled by default (increment 0), so behavior is unchanged unless opted in.

How was this patch tested?

New unit tests in TaskSetManagerSuite and TaskSchedulerImplSuite covering the disabled default (zero regression), cpus increment, strict wait, no starvation of normal tasks, executor-cores cap, reset on success, both OOM recognition paths, configurable exit codes, barrier exclusion, speculative consistency, and the maxTaskFailures safety bound.

[sunchao]

Summary

Prior state and problem

Spark currently retries an OOM-failed task with the ResourceProfile's original CPU allocation, so a retry may repeat the same memory-pressure conditions. This change introduces an opt-in mechanism that increases CPUs on successive OOM retries to reduce executor concurrency and give the retry a larger execution-memory share.

Design approach

The patch keeps a per-partition OOM counter, recognizes SparkOutOfMemoryError and configured executor exit codes, and calculates baseCpus + increment * retries. It caps the result at the executor core count and skips an enlarged retry when the current offer lacks enough free CPUs.

Correctness / compatibility analysis

The disabled-default path appears unchanged, and the coarse-grained backend correctly charges and releases TaskDescription.cpus. Barrier exclusion, speculative-attempt sharing, success reset, custom-resource handling, and the ordinary maximum-failure bound also look consistent.

The enabled path still has correctness gaps in executor-core discovery, fatal-OOM event ordering, capacity acquisition, local-mode accounting, and integer arithmetic. These can disable the feature, violate the base CPU request, indefinitely delay an OOM retry, oversubscribe local mode, or leave scheduler bookkeeping inconsistent.

Key design decisions

Using a per-task-index counter and resetting it after success are reasonable choices. The risky decisions are interpreting a missing executor-core value as one, recognizing a JVM OOM only during executor loss, and implementing strict wait without reserving capacity across offers.

Implementation sketch

The patch adds two configurations and changes TaskSetManager to compute a variable CPU request for each retry. TaskSchedulerImpl then accounts for the CPUs in each returned TaskDescription, with unit tests and user-facing documentation for the new behavior.

Behavioral changes worth calling out

Although the feature is disabled by default, enabling it changes task resource requirements dynamically, which affects offer selection, fairness, dynamic allocation, and scheduler-backend accounting. Those downstream paths do not all understand the enlarged task shape yet.

Suggested improvements

Please address the inline findings and add deterministic coverage for missing executor-core metadata, FAILED-before-executor-loss ordering, repeated partial offers, real local-backend accounting, and the maximum accepted increment value. The core CI portion passed; the aggregate build failure is from Minikube's Kubernetes-version lookup hitting GitHub API rate limits before the Kubernetes tests began.

[sunchao]

[P1] Preserve the ResourceProfile's base CPU request when executor cores are unknown

For Standalone and local-cluster without an explicit spark.executor.cores, getExecutorCores is intentionally empty because the executor can use the worker's available cores, while conf.get(EXECUTOR_CORES) returns the generic default of 1. Once the increment is enabled, effectiveCpusFor is called for every attempt, including when numOomRetries(index) == 0, so a base spark.task.cpus=2 is reduced to one CPU. TaskSchedulerImpl prechecks the offer against the two-CPU ResourceProfile request but later debits the one CPU in TaskDescription, so an eight-core executor can launch seven nominally two-CPU tasks instead of four; with the default one CPU per task, OOM retries never grow at all. Please derive this cap from the executor's actual capacity (or represent an unknown cap explicitly) and ensure effectiveCpusFor can never return less than baseCpus.

[ulysses-you]

Fixed in 8f657a7. executorCoresLimit is now an Option[Int]: it takes the ResourceProfile's executor cores, falls back to an explicitly configured spark.executor.cores, and is None only when neither is set (Standalone / local-cluster using the worker's cores) rather than collapsing to the generic default of 1. effectiveCpusFor applies the cap only when known and is floored at baseCpus, so it can never return less than the base request. Added a test (OOM retry grows cpus when executor cores are unknown) that builds a TaskResourceProfile with getExecutorCores == None and asserts the retry still grows from the base instead of being pinned at 1.

[sunchao]

[P1] Reserve capacity for enlarged retries across offer rounds

This condition leaves an enlarged retry queued when the current offer is too small, but then continues scanning and may launch a normal one-CPU task. On a saturated executor, CoarseGrainedSchedulerBackend makes a new offer immediately after each individual task completion; each released core can therefore be backfilled before two or more cores ever accumulate. The retry is delayed until the normal backlog drains, and continuing work from competing task sets can keep it waiting indefinitely. Dynamic allocation also derives executor demand from the ResourceProfile's base maxTasksPerExecutor, so it need not add an executor for this larger pending task. Please add a reservation/draining mechanism across offer rounds and make allocation aware of the enlarged task shape.

[sunchao]

[P2] Compute retry CPU growth without Int overflow

The new configuration accepts every nonnegative Int, so with baseCpus = 1 and oomRetryCpusIncrement = Int.MaxValue, the first OOM retry wraps requested to Int.MinValue before this clamp. oomRetryNeedsMoreCpus then accepts the negative value, and prepareLaunchingTask mutates copiesRunning, taskInfos, and the running-task set before TaskDescription asserts that cpus > 0. The launch fails with scheduler state already recording a running copy. Please calculate in Long and clamp safely before converting to Int, or impose a configuration bound that makes the arithmetic provably safe.

[ulysses-you]

Fixed in 8f657a7. effectiveCpusFor now floors the result at baseCpus:

private def effectiveCpusFor(index: Int, baseCpus: Int): Int = {
  val requested = baseCpus + oomRetryCpusIncrement * numOomRetries(index)
  executorCoresLimit.map(_.min(requested)).getOrElse(requested).max(baseCpus)
}

Even when an extreme increment overflows the Int arithmetic to a negative requested, the .max(baseCpus) floor guarantees a positive result, so oomRetryNeedsMoreCpus and prepareLaunchingTask never see a non-positive value and the cpus > 0 assertion in TaskDescription is never reached with scheduler state already mutated. The request degrades gracefully to baseCpus in that case. Added a test (maximum increment never yields a non-positive cpu request) that OOM-fails with oomRetryCpusIncrement = Int.MaxValue and asserts the retry still launches with a positive cpu count.

[sunchao]

[P1] Preserve JVM-OOM classification when FAILED arrives before executor loss

Executor.TaskRunner sends an ExceptionFailure status before invoking the fatal exception handler that exits with code 52. If the driver processes that FAILED update first, handleFailedTask marks the attempt finished; its OOM branch recognizes only SparkOutOfMemoryError, not java.lang.OutOfMemoryError. When the subsequent ExecutorExited(52, ...) reaches this block, info.running is already false, so the counter is never incremented and the retry keeps its original CPU count. This is a normal reachable ordering for the advertised JVM-heap-OOM path. Please classify fatal OutOfMemoryError in the ExceptionFailure path or preserve enough attempt state for the later definitive loss reason, and add a test with FAILED preceding executor loss.

[ulysses-you]

Fixed in 8f657a7. The ExceptionFailure path now recognizes any OutOfMemoryError, not only SparkOutOfMemoryError, via a helper:

private def isOom(ef: ExceptionFailure): Boolean = {
  ef.exception.exists(_.isInstanceOf[OutOfMemoryError]) ||
    ef.className == classOf[OutOfMemoryError].getName ||
    ef.className == classOf[SparkOutOfMemoryError].getName
}

It prefers the preserved Throwable and falls back to the class name, since the exception may not be preserved (preserveCause = false) or loadable in the driver. So a fatal java.lang.OutOfMemoryError reported through the ExceptionFailure (which the Executor sends before it exits with code 52) now grows the retry's cpus even when the FAILED update is processed before the executor-loss event. Double counting is avoided in both orderings: whichever event arrives first does the increment, and the info.finished guard at the top of handleFailedTask makes the second a no-op. Added two tests covering FAILED-before-executor-loss and the reverse ordering, each asserting the counter increments exactly once.

[sunchao]

[P2] Keep LocalSchedulerBackend free-core accounting in sync

This scheduler-side change correctly charges TaskDescription.cpus, but LocalSchedulerBackend still subtracts and restores the fixed scheduler.CPUS_PER_TASK. For example, in local[4] with spark.executor.cores=4, base CPUs 1, and a two-CPU retry, launching that retry plus two normal tasks consumes four CPUs here while the local endpoint subtracts only three and advertises a phantom core. Another revive or task-set submission can then oversubscribe the local executor. Please track each launched task's actual CPU count in the local backend and restore that count on completion.

[ulysses-you]

Fixed in 8f657a7. LocalEndpoint now tracks each launched task's actual cpus in a taskId -> cpus map and refunds that exact count on completion, instead of the fixed CPUS_PER_TASK:

private val runningTaskCpus = new mutable.HashMap[Long, Int]()
...
// StatusUpdate, on finished:
freeCores += runningTaskCpus.remove(taskId).getOrElse(scheduler.CPUS_PER_TASK)
...
// reviveOffers, per launched task:
freeCores -= task.cpus
runningTaskCpus(task.taskId) = task.cpus

Charging task.cpus at launch keeps the local free-core count consistent with what TaskSchedulerImpl debited, so a two-cpu OOM retry no longer leaves a phantom core that could oversubscribe the local executor. Added an end-to-end test in FailureSuite that runs a local[4,2] job whose first attempts OOM and whose retries need more cpus, asserting the job completes with the expected task count.