[SPARK-49485][CORE] Request additional executor for speculative tasks when active executors equal maxNeeded

[zahed1994]

What changes were proposed in this pull request?

When spark.dynamicAllocation.enabled=true and spark.speculation=true, straggling tasks requiring speculative execution can stall indefinitely if all remaining active executors reside on the same host.

In ExecutorAllocationManager.scala, maxNumExecutorsNeededPerResourceProfile calculates maxNeeded strictly based on (running + pendingTasks + pendingSpeculative) / tasksPerExecutor. If the current active executor count equals maxNeeded, ExecutorAllocationManager calculates target executors as equal to the active count and requests no new executors. At the same time, Spark's task scheduler avoids launching speculative task copies on an executor on the same host where the task is already running slow. Consequently, no active executor can run the speculative task and no new executor is requested, causing speculative tasks to stall indefinitely.

This PR updates maxNumExecutorsNeededPerResourceProfile in ExecutorAllocationManager.scala to allocate an additional target executor when pendingSpeculative > 0 and maxNeeded equals the current active executor count. This allows ExecutorAllocationManager to request an extra executor from the cluster manager (YARN / K8s / Standalone) on a distinct host to execute the speculative task.

Why are the changes needed?

Without this change, applications using Dynamic Allocation and speculation can hang indefinitely when remaining executors reside on the same slow host.

Does this PR introduce any user-facing change?

No.

How was this patch tested?

• Added unit test SPARK-49485: request additional executor when speculative tasks equal maxNeeded in ExecutorAllocationManagerSuite.scala.
• Verified cleanly via core/compile and core/scalastyle.

[sunchao]

[P2] Exercise the new allocation condition in the regression test

createConf(1, 5, 2) sets the initial executor count, leaving one task slot per executor. After speculation is submitted, maxNeeded is already ceil((2 + 1) / 1) = 3, while the active count is 2. The new branch never executes, so this assertion cannot detect removal of the fix.

For example, configure two two-core executors with three running tasks and one pending speculative task: base then returns 2 and head returns 3.

[zahed1994]

Thanks for the review @sunchao! Good catch , with tasksPerExecutor = 1, ceil((2 + 1) / 1) = 3 was already returning 3 on master even without the fix.

I have updated the regression test to configure 2-core executors (spark.executor.cores = 2) with 3 running tasks across 2 active executors and 1 pending speculative task (4 total tasks).

Without this fix (base), maxNeeded evaluates to ceil(4 / 2) = 2. With this fix (head), since maxNeeded (2) equals the active executor count (2) and pendingSpeculative > 0, the new offset triggers and requests 2 + 1 = 3 executors, properly exercising the new allocation branch.

[sunchao]

[P2] Remove the extra pending regular task from the fixture

The revised fixture still leaves one regular task pending: createStageInfo(0, 4) declares four regular tasks, but only three TaskStart events are posted. After speculative submission, raw maxNeeded is ceil((3 running + 1 regular pending + 1 speculative pending) / 2) = 3, while the executor count is 2. The new branch still never executes, so the assertion passes without the fix.

Use createStageInfo(0, 3) to establish the intended base=2/head=3 distinction.

[sunchao]

[P2] Wrap the comment to restore Scala lint

This revised comment is 101 characters long, exceeding the configured 100-character limit. The current CI Scala-linter step fails at this exact line with File line length exceeds 100 characters. Wrap or shorten the comment.

[zahed1994]

Thanks for the review @sunchao! I've updated the PR accordingly:

1. Fixture stage task count: Changed createStageInfo(0, 4) to createStageInfo(0, 3) so there are zero regular pending tasks (3 running + 0 pending + 1 speculative = 4 total tasks). Without this fix, maxNeeded evaluates to ceil(4 / 2) = 2, whereas with this fix, 2 + 1 = 3 triggers, strictly exercising the new allocation path (base=2 / head=3).
2. Scala Linter: Shortened and wrapped the comment lines so all lines are strictly under 100 characters (max line length 94 chars).

All tests and linters are passing cleanly now. Could you please take another look when you get a chance?

[sunchao]

[P2] Preserve the extra executor with the StatefulSet allocator

At reviewed commit 9b194f91389bed3cdd5ace9509bb2d67ef5c4085, the new condition in ExecutorAllocationManager.scala can introduce executor allocation/deletion churn when dynamic allocation and speculation are enabled with spark.kubernetes.allocation.pods.allocator=statefulset.

For example, with two 2-core executors, allocation ratio 1, and a maximum of at least 3 executors, one executor can be full while the other runs a slow task whose speculative copy cannot use the spare slot on the original host. Three running tasks plus one pending speculative copy gives ceil(4 / 2) = 2, so this branch raises the requested count to 3. Once the third executor registers, the equality check no longer holds and the next allocation update reduces the request to 2. Starting the copy does not prevent that reduction: four running tasks still give a calculated requirement of 2.

StatefulSetPodsAllocator directly applies that reduced request with scale(expected, false). StatefulSet scale-down can therefore delete the newly added highest-ordinal pod, potentially discarding the speculative attempt. Allocation/deletion can repeat while the original remains slow and no other eligible slot becomes available.

The allocator's immediate downscale behavior predates this PR; the newly introduced part is the 2 -> 3 -> 2 request at unchanged task demand, where the base stays at 2. This concern is specific to the StatefulSet allocator, not the default direct allocator. Could we handle or guard this combination so the additional executor survives long enough to do useful speculative work?

This is based on source inspection, not a local Kubernetes runtime reproduction. The earlier fixture and lint issues are fixed.

[zahed1994]

Thanks @sunchao, good catch on the StatefulSet case.

You're right that the previous logic could cause the target to drop back to 2 once the additional executor was registered or the speculative task moved from pending to running. In the StatefulSet allocator, that can immediately trigger a scale-down and remove the newly allocated executor before the speculative task has a chance to make progress.

I've updated the allocation logic in ExecutorAllocationManager.scala to account for the full lifecycle of the speculative task rather than only the pending state.

Specifically:

1. Track the full speculative task lifecycle

   ExecutorAllocationListener now provides the number of speculative tasks that are either pending or currently running. The allocation logic uses:

   speculativeTasks = pendingSpeculative + runningSpeculative

   This ensures that the additional executor remains accounted for after the speculative task has started running.

2. Use the regular-task requirement as the baseline

   Instead of comparing maxNeeded against the current active executor count, I calculate baseMaxNeededWithoutSpeculation from the regular pending/running tasks (excluding running speculative tasks).

   The locality offset is then applied when:

   maxNeeded == baseMaxNeededWithoutSpeculation

   and there is an active speculative task. This keeps the additional executor tied to the actual resource requirement of the regular workload, rather than the transient number of executors that happen to be registered at that point.

3. Retain the additional executor while speculation is running

   With this change, the allocation remains stable across the speculative task transition:

   pending → running

   So in the case you described, the target remains at 3 after executor SPARK-1135: fix broken anchors in docs #3 registers and after the speculative task starts, instead of immediately returning to 2 and allowing the StatefulSet allocator to remove the additional executor.

   Once the speculative attempt finishes, the speculative-task count goes back to zero and the target can return to the regular-task requirement of 2.

4. Handle the killed-speculation path

   I also added coverage for the case where the speculative attempt is killed rather than completing normally. The additional executor is retained while the speculative task is running and released once the speculative attempt ends.

I've added regression coverage in ExecutorAllocationManagerSuite for both:

pending → running → completed

and

pending → running → killed

The targeted tests and linters are passing cleanly.

Could you please take another look when you get a chance?

[sunchao]

The previously reported StatefulSet registration issue is addressed in reviewed commit 1ac32efce4cb93c6ea31d6955af4d54f09ceede6. Two P2 issues remain.

[P2] Count running speculative attempts, not distinct task indices

The new getRunningSpeculativeTaskSum returns the size of a set keyed by task index. After a primary attempt fails, its surviving speculative attempt S1 can itself be speculated: TaskSetManager checks every running attempt, and the dequeue guard permits another speculative attempt when copiesRunning(index) == 1. S1 and S2 both have speculative=true, but their starts produce only one set entry. Ending either attempt also removes that entry while the other remains running.

A concrete allocation consequence uses three task slots per executor and allocation ratio 1. Let a two-task result stage have one completed fast task and a slow primary P; two unrelated regular tasks from another stage use the same resource profile. Executor B runs those two regular tasks plus S1, the speculative copy of P. P's executor A fails, leaving one regular task pending, and replacement executor C can later launch S2 after S1 becomes speculatable again.

While S2 is pending, the estimate is 2. Once S2 starts, there are still five running-or-pending attempts: one regular pending and four running, of which two are speculative. The accessor reports only one running speculative attempt, so the regular baseline becomes ceil((1 + 4 - 1) / 3) = 2. Since raw maxNeeded is also 2, the new branch raises the target to 3. Counting both speculative attempts gives a regular baseline of 1 and keeps the target at 2, as do the base and previous revision after launch. The regular pending entry keeps the backlog timer active, so this can request an unnecessary executor even though B and C already run all four attempts.

Please track live speculative attempt IDs or equivalent multiplicities, and cover re-speculation after a primary failure. This finding is based on source inspection, not a local runtime reproduction.

[P2] Update the existing regression expectations for retained speculation

The existing SPARK-41192 assertions at lines 758-759 and SPARK-30511 assertions at lines 865-866 still expect 1 executor while speculative attempts are running. The revised retention behavior intentionally returns 2 in both cases, and the current core CI job fails both tests with 2 did not equal 1.

The later assertions at lines 886-887 also need the corresponding update; that mismatch is source-established because the earlier failure prevents CI from reaching it. These failures call for updating the old expectations, not reverting the intended retention behavior.

Validation: the exact-head CI artifact reports 47 passes and 2 failures out of 49 allocation tests. Both new tests pass, and Scala lint passes.

[zahed1994]

Thanks @sunchao, good catches on both points.

I've updated the PR in the latest commit:

1. Track running speculative attempts by taskId: Changed
   stageAttemptToSpeculativeTaskIndices to
   stageAttemptToSpeculativeTaskIds, storing taskInfo.taskId as a Long
   instead of the logical task index. This ensures multiple speculative
   attempts for the same task index are counted independently. Added a
   regression test covering multiple running speculative attempts for the
   same task index.

2. Update retention expectations: Updated the existing SPARK-41192 and
   SPARK-30511 assertions to expect 2 executors while speculative attempts
   remain active, reflecting the intended retention behavior.

The ExecutorAllocationManagerSuite passes all 50 tests, including the new
regression coverage, and scalastyle passes cleanly.

Could you please take another look when you get a chance?