[SPARK-57491][FOLLOWUP] Make stale push-based shuffle fallback chunk-granular and opt-in by gaoyajun02 · Pull Request #58008 · apache/spark

gaoyajun02 · 2026-08-14T10:03:53Z

What changes were proposed in this pull request?

Follow-up to SPARK-57491 (stale push-based shuffle data detection). SPARK-57491's reducer fallback was too coarse; this PR makes it chunk-granular and opt-in.

Chunk-granularity fallback. Replaces checkStaleMapIdInMergedBlock (which fell back the entire merged block if any stale map index was present) with isStaleChunk, which checks a single chunk's bitmap. The meta fetch no longer short-circuits, so in ShuffleBlockFetcherIterator.next() only the chunk that actually contains a stale index falls back to its original blocks; other chunks of the same merged block keep being read from the merged block.
No cascade for stale fallback. initiateFallbackFetchForPushMergedBlock gains a fallbackPendingChunks flag (default true). Fetch-failure paths keep cascading to same-host pending chunks; the stale path passes false, since a stale chunk doesn't imply the host is unhealthy, so pending chunks retain their merged-block reads.
Opt-in, indeterminate-scoped reducer fallback. Two new configs (both default false): spark.shuffle.push.stale.fallback.enabled gates markStalePushedPartition; spark.shuffle.push.stale.detectAllStages.enabled extends it to all map stages. Duplicate map attempts are now always logged for observability, but marked stale only when fallback is enabled and the stage is in scope. Stage indeterminacy is resolved once at TaskSetManager construction (isIndeterminateShuffleMapStage).

Why are the changes needed?

Speculation is common on large stages in production, so duplicate map attempts that both push are routine. SPARK-57491's fallback had too wide a blast radius:

Whole-block fallback amplifies fetch load across the whole stage. A stale index in one chunk forced fallback of the entire merged block, turning 1 merged fetch into N original-block fetches per affected reducer. On a large shuffle, a single stale partition multiplies network requests by an order of magnitude and slows the whole stage, rather than protecting just the one partition at risk.
Deterministic stages pay for nothing. Duplicate deterministic attempts produce byte-identical output, so the data in the merged block is actually correct — the fallback is pure overhead and a false positive. The real risk is confined to indeterminate stages, where duplicate attempts can diverge.
Safe rollout. SPARK-57491 is unreleased; its always-on fallback could land as a silent perf regression for any job that hits speculation. Opt-in (default: log-only) fallback lets the detection ship safely, enabled only where the correctness gain justifies the extra fetches.

Net effect: a stale partition now costs roughly 1 chunk's worth of extra fetches instead of a full merged block, and only indeterminate stages pay it.

Does this PR introduce any user-facing change?

Yes — two new optional configs (both default false): spark.shuffle.push.stale.fallback.enabled and spark.shuffle.push.stale.detectAllStages.enabled.

How was this patch tested?

Added/updated unit tests:

TaskSetManagerSuite: updated the SPARK-57491 test to enable both fallback switches and assert the partition is marked stale; new test verifies that with fallback disabled (default) the stale attempt is reported but not marked.
ShuffleBlockFetcherIteratorSuite: updated to chunk granularity — a stale index in a chunk's bitmap triggers fallback for that chunk only; new case asserts the non-stale chunk is still read as a ShuffleBlockChunkId.

Was this patch authored or co-authored using generative AI tooling?

Yes. Generated-by: Claude Code (Anthropic)

gaoyajun02 · 2026-08-14T10:11:17Z

PTAL @cloud-fan

…granular and opt-in

cloud-fan

2 blocking, 2 non-blocking, 2 nits.
Changes requested: the chunk-level fallback shape is sound, but the rollout default and runtime-indeterminacy lifecycle need correction before merge.

Design / architecture (2)

Blocking: core/src/main/scala/org/apache/spark/internal/config/package.scala:2985: Make stale-push fallback actually opt-in by changing the registered default to false. -- see inline
Blocking: core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala:89: Evaluate runtime indeterminacy when the duplicate arrives instead of caching it at TaskSetManager construction. -- see inline

Correctness (2)

Non-blocking: core/src/main/scala/org/apache/spark/storage/PushBasedFetchHelper.scala:301: Avoid copying the entire stale-index set for every fetched chunk. -- see inline
Non-blocking: core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala:2930: Test the real default and indeterminate-only mode independently. -- see inline

Nits: 2 minor items (see inline comments).

Verification

Traced a late duplicate from TaskSetManager.handleSuccessfulTask through stale marking and MapOutputTracker propagation to the reducer's chunk predicate and single-chunk fallback. Confirmed that DAGScheduler can set runtime indeterminacy after TaskSetManager construction, and that the worker getter copies the complete stale-index set on each chunk check.

cloud-fan · 2026-08-20T16:10:11Z

+ "spark.shuffle.push.stale.detectAllStages.enabled is also true.")
+ .version("4.4.0")
+ .booleanConf
+ .createWithDefault(true)


Blocking:
The rollout contract is opt-in: both this entry's documentation and the PR description say the default is false. Leaving this as true enables fallback for indeterminate stages without an explicit setting and reintroduces the fetch amplification this PR is meant to avoid.
Suggested change
.createWithDefault(true)
.createWithDefault(false)

Done. Changed createWithDefault(true) -> createWithDefault(false) so the registered default matches the "Default false" documentation and the opt-in contract. The default no longer re-enables reducer fallback for indeterminate stages without an explicit setting.
Verified the only reader is TaskSetManager.stalePushFallbackEnabled; tests set the config explicitly, so no test relied on the true default.

cloud-fan · 2026-08-20T16:10:11Z

+ // identical output across attempts, so a stale push there is benign). Defaults to false when the
+ // stage isn't registered (e.g. tests). Uses taskSet.stageId (constructor param) because the
+ // `stageId` field below is not initialized yet at this source position.
+ private val isIndeterminateShuffleMapStage: Boolean = isShuffleMapTasks &&


Blocking:
Keep the ShuffleMapStage reference and evaluate its indeterminacy when the duplicate result arrives, rather than caching this Boolean. DAGScheduler can set isChecksumMismatched later in the same stage attempt; a subsequent duplicate then still sees false here, skips stale marking in the indeterminate-only mode, and leaves reducers able to consume the stale merged chunk.

Done. TaskSetManager no longer caches indeterminacy as a Boolean at construction. It now holds the Option[ShuffleMapStage] reference and re-evaluates isStaticallyIndeterminate || isRuntimeIndeterminate inside detectStalePushIfShuffleTask, when the late duplicate result arrives. A checksum mismatch set by DAGScheduler during the stage (via registerMapOutput) is therefore observed rather than read from a stale constructor snapshot.
Added a test case (fallbackEnabled=true, detectAllStages=false) that sets isChecksumMismatched after the TaskSetManager is constructed and asserts stale marking still fires — this would fail under the old cached-Boolean approach.

cloud-fan · 2026-08-20T16:10:11Z

 // shuffleId is only available when isShuffleMapTasks=true
 private val shuffleId = taskSet.shuffleId
+ // Scopes stale-push reducer fallback to indeterminate stages (deterministic stages reproduce
+ // identical output across attempts, so a stale push there is benign). Defaults to false when the


Nit:
Please phrase this as reproducing the same data set. Spark's UNORDERED level is not indeterminate, but its contract explicitly permits record order to change across reruns, so identical output is stronger than the invariant used by this guard.

Done. Reworded to "reproduce the same data set" — UNORDERED is not indeterminate but permits record reordering, so "identical output" was stronger than the invariant this guard relies on.

cloud-fan · 2026-08-20T16:10:11Z

+ * the stale pushed map indexes recorded by the MapOutputTracker, which are non-empty only
+ * when reducer fallback marking is enabled.
+ *
+ * Each invocation gets its own SparkContext via [[withSpark]] so the caller can loop without


Nit:
The unqualified member does not resolve from this suite.
Suggested change
*Each invocation gets its own SparkContext via [[withSpark]] so the caller can loop without
*Each invocation gets its own SparkContext via [[LocalSparkContext.withSpark]] so the caller can loop without

Done. Changed [[withSpark]] -> [[LocalSparkContext.withSpark]].

cloud-fan · 2026-08-20T16:10:11Z

- logWarning(s"Found stale pushed map indexes in merged block $shuffleBlockId from" +
- s" ${address.host}:${address.port}, falling back to fetch the original blocks")
+ private[spark] def isStaleChunk(blockId: ShuffleBlockChunkId): Boolean = {
+ val staleMapIndexes = mapOutputTracker.getStaleMapIndexes(blockId.shuffleId)


Non-blocking:
This runs once per fetched chunk, while MapOutputTrackerWorker.getStaleMapIndexes allocates a new HashSet and copies every stale index on each call. A partition with C chunks and S stale maps therefore copies O(C * S) entries in the reducer fetch path. Please move the intersection behind a tracker query that inspects the published snapshot without returning a defensive copy.

Done. Added intersectsStaleMapIndexes(shuffleId, bitmap) to MapOutputTracker (trait + master + worker) and ShuffleStatus, which inspects the published stale snapshot via iterator + bitmap.contains without returning a defensive copy. isStaleChunk now delegates to it, dropping the O(C * S) copy on the reducer fetch path.
getStaleMapIndexes is kept for the driver-side propagation serialization path, where a copy is still appropriate.

cloud-fan · 2026-08-20T16:10:11Z

+ // Both switches enabled: the late attempt is marked stale regardless of stage determinism
+ // (the test does not register an indeterminate ShuffleMapStage, so detectAllStages is
+ // required here).
+ (true, true, true),


Non-blocking:
This matrix bypasses both contracts it needs to protect: the disabled row sets false explicitly, and the enabled row turns on detectAllStages. Leave the fallback setting unset in one case, and add an enabled/detect-all-disabled case with a registered stage whose runtime indeterminacy changes after manager construction; those cases catch the current default mismatch and stale constructor snapshot.

Done. The matrix now takes Option[Boolean] so either switch can be left unset:
(None, None, expectStale=false) — exercises the real default (no explicit config set); catches the default mismatch.
(Some(true), Some(false), expectStale=true, runtimeIndeterminateAfterConstruction=true) — enabled / detect-all-disabled, with a registered ShuffleMapStage whose isChecksumMismatched is set afterTaskSetManager construction; catches the cached-snapshot issue (the indeterminate-only path).
The helper registers a real ShuffleMapStage in dagScheduler.stageIdToStage so the reference resolves, matching production.

…erminacy, stale-set lifecycle Address cloud-fan's review (2 blocking, 2 non-blocking, 2 nits) on apache#58008 plus a related stale-set lifecycle bug found during review. Blocking: - spark.shuffle.push.stale.fallback.enabled registered default true -> false, matching its "Default false" doc and the opt-in contract (default true re-enabled fetch amplification for indeterminate stages). - TaskSetManager no longer caches stage indeterminacy as a Boolean at construction. It now holds the ShuffleMapStage reference and re-evaluates isStaticallyIndeterminate || isRuntimeIndeterminate when a late duplicate result arrives, so a checksum mismatch set by DAGScheduler during the stage is observed. Non-blocking: - isStaleChunk delegates to a new intersectsStaleMapIndexes(shuffleId, bitmap) tracker query that inspects the published stale snapshot without returning a defensive copy per fetched chunk (was O(C * S) copies on the reducer fetch path). - Test matrix now covers the real default (config unset) and an enabled/detect-all-disabled case with runtime indeterminacy toggled after TaskSetManager construction, catching both the default mismatch and the cached indeterminacy snapshot. Related fix: - unregisterAllMapAndMergeOutput now clears ShuffleStatus.staleMapIndexes. Without this, stale marks from a previous attempt survived a stage retry and caused spurious chunk fallback on the retried attempt's fresh push data. The worker-side stale cache is refreshed via the existing epoch bump. Nits: - "identical output" -> "the same data set" (UNORDERED permits reordering). - [[withSpark]] -> [[LocalSparkContext.withSpark]] (unqualified member did not resolve). Co-Authored-By: Claude <noreply@anthropic.com>

gaoyajun02 · 2026-08-25T09:40:00Z

Addressed all 6 review items in 1eddc72 (2 blocking, 2 non-blocking, 2 nits), plus a related stale-set lifecycle bug surfaced while investigating #2:

Removed reference to incubation in README.md. #1 (blocking):fallback.enabled registered default true -> false.
Removed reference to incubation in Spark user docs. #2 (blocking):TaskSetManager holds the ShuffleMapStage reference and re-evaluates indeterminacy when the late duplicate arrives (was cached as a Boolean at construction).
SPARK-1135: fix broken anchors in docs #3 (non-blocking):isStaleChunk delegates to a new intersectsStaleMapIndexes tracker query — no per-chunk defensive copy.
SPARK-1137: Make ZK PersistenceEngine not crash for wrong serialVersionUID #4 (non-blocking): test matrix now covers the real default (config unset) and an enabled/detect-all-disabled case with runtime indeterminacy toggled post-construction.
Nits: "identical output" -> "the same data set"; [[withSpark]] -> [[LocalSparkContext.withSpark]].
Related:unregisterAllMapAndMergeOutput now clears staleMapIndexes — without this, stale marks from a previous attempt survived a stage retry and caused spurious chunk fallback on the retried attempt's fresh push data. Added an assertion to the existing propagation test.

core/compile, core/Test/compile, and the full TaskSetManagerSuite / MapOutputTrackerSuite / ShuffleBlockFetcherIteratorSuite pass.

PTAL @cloud-fan

gaoyajun02 changed the title ~~[SPARK-33235][FOLLOWUP] Make stale push-based shuffle fallback chunk-granular and opt-in~~[SPARK-57491][FOLLOWUP] Make stale push-based shuffle fallback chunk-granular and opt-inAug 14, 2026

gaoyajun02 force-pushed the SPARK-33235-fixup branch from a0786f5 to a36c3b8Compare August 14, 2026 10:08

gaoyajun02 force-pushed the SPARK-33235-fixup branch from a36c3b8 to 52d6c01Compare August 14, 2026 13:22

[SPARK-57491][FOLLOWUP] Make stale push-based shuffle fallback chunk-…
091123c
…granular and opt-in

gaoyajun02 force-pushed the SPARK-33235-fixup branch from 52d6c01 to 091123cCompare August 20, 2026 05:35

cloud-fan reviewed Aug 20, 2026
View reviewed changes

	*Each invocation gets its own SparkContext via [[withSpark]] so the caller can loop without
	*Each invocation gets its own SparkContext via [[LocalSparkContext.withSpark]] so the caller can loop without

Uh oh!

Conversation

gaoyajun02 commented Aug 14, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

Was this patch authored or co-authored using generative AI tooling?

Uh oh!

gaoyajun02 commented Aug 14, 2026

Uh oh!

cloud-fan left a comment

Choose a reason for hiding this comment

Design / architecture (2)

Correctness (2)

Verification

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gaoyajun02 commented Aug 25, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

gaoyajun02 commented Aug 14, 2026•
edited
Loading