This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 970b4dc006eb [SPARK-57491][CORE] Prevent and detect stale push-based
shuffle data from duplicate map task attempts
970b4dc006eb is described below
commit 970b4dc006ebcc1a5085c45c970340861ddc2f9c
Author: gaoyajun02 <[email protected]>
AuthorDate: Thu Jul 9 00:48:39 2026 +0800
[SPARK-57491][CORE] Prevent and detect stale push-based shuffle data from
duplicate map task attempts
### What changes were proposed in this pull request?
Fix push-based shuffle serving stale/incorrect data when multiple task
attempts for the same map partition both push data to the external merger.
Three layers of defense:
- **ShuffleWriteProcessor** — Defer shuffle block push via
`PostStatusUpdateListener` so that killed/failed tasks never push. The listener
checks `context.isInterrupted()` / `context.isFailed()` before initiating push,
and logs which task was skipped and why.
- **TaskSetManager + MapOutputTracker** — Track stale (duplicate)
partitionIds on the driver side. When a speculative or retried ShuffleMapTask
result arrives after another attempt for the same partition has already
committed, `TaskSetManager.detectStalePushIfShuffleTask` marks the partition's
ID as stale in `ShuffleStatus.staleMapIndexes`. The stale set is serialized
alongside existing map/merge output metadata and propagated to executor-side
`MapOutputTrackerWorker`.
- **PushBasedFetchHelper (`ShuffleBlockFetcherIterator`)** — Before reading
any merged block on the reducer side, check chunk-level RoaringBitmaps for
stale partitionIds via the new `checkStaleMapIdInMergedBlock` method. If a
stale mapIndex is found in any chunk bitmap, log a WARN and fall back to
fetching original unmerged blocks. When the stale set is empty (the common
case), this check short-circuits immediately with zero overhead.
### Why are the changes needed?
With push-based shuffle enabled, indeterminate stages can produce incorrect
results when speculation or task retry causes two attempts of the same map
partition to both push data to the external merger. The merged block ends up
with interleaved data from both attempts, but only one attempt's MapStatus is
committed on the driver. Downstream reducers silently read corrupted/duplicated
data.
This is triggered by non-deterministic functions (`rand()`, `UUID()`, etc.)
in shuffle keys or join conditions where row ordering differs between attempts.
Even with deferred push (layer 1), there is a race window: if an original
attempt succeeds and pushes, then gets killed by a speculative winner that also
pushes, the merger receives data from both attempts but only one MapStatus is
registered.
Why the existing checksum-mismatch / rollback machinery does not cover this
case？
The existing checksum-mismatch detection in
`DAGScheduler.handleTaskCompletion` relies on comparing
`MapStatus.checksumValue` values via `MapOutputTracker.registerMapOutput()`.
This only works for task results reported as **Success** — the `MapStatus` is
registered and compared against any previously registered status for the same
partition.
In speculation with push-based shuffle, the **losing attempt's result never
reaches `registerMapOutput()`**: when a later-arriving or killed speculative
attempt is detected in `TaskSetManager.handleSuccessfulTask`, it is handled as
`TaskKilled` (not `Success`) and forwarded to DAGScheduler as a kill event,
which bypasses the `registerMapOutput` path entirely. Therefore no checksum
comparison occurs for the loser's output.
Furthermore, even if this path were reachable,
`checksumMismatchFullRetryEnabled` is **force-disabled** when push-based
shuffle is active (`Dependency.scala:148-149`), since the two features are
designed for mutually exclusive scenarios: checksum mismatch triggers
stage-level rollback with full retry of all downstream stages, while stale push
detection handles the problem at the partition level with a lightweight fetch
fallback.
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
- UT
- Manual verification on internal cluster confirmed: speculative duplicate
pushes are either prevented by deferred push or detected and fallen back from,
with expected WARN log on fallback.
### Was this patch authored or co-authored using generative AI tooling?
Yes, co-authored with GLM-5V-Turbo.
Closes #56559 from gaoyajun02/SPARK-33235.
Lead-authored-by: gaoyajun02 <[email protected]>
Co-authored-by: gaoyajun02 <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
(cherry picked from commit 50932d2f0f1aca02d594f4d8bc367a5ad91cf840)
Signed-off-by: Wenchen Fan <[email protected]>
---
.../org/apache/spark/BarrierTaskContext.scala | 9 ++
.../scala/org/apache/spark/MapOutputTracker.scala | 149 ++++++++++++++++++++-
.../main/scala/org/apache/spark/TaskContext.scala | 19 ++-
.../scala/org/apache/spark/TaskContextImpl.scala | 20 +++
.../scala/org/apache/spark/executor/Executor.scala | 3 +
.../apache/spark/scheduler/TaskSetManager.scala | 30 +++++
.../spark/shuffle/ShuffleWriteProcessor.scala | 26 +++-
.../spark/storage/PushBasedFetchHelper.scala | 54 +++++++-
.../org/apache/spark/util/taskListeners.scala | 17 +++
.../org/apache/spark/MapOutputTrackerSuite.scala | 54 +++++++-
.../apache/spark/scheduler/TaskContextSuite.scala | 70 ++++++++++
.../spark/scheduler/TaskSetManagerSuite.scala | 83 ++++++++++++
.../storage/ShuffleBlockFetcherIteratorSuite.scala | 81 +++++++++++
project/MimaExcludes.scala | 5 +-
14 files changed, 603 insertions(+), 17 deletions(-)
diff --git a/core/src/main/scala/org/apache/spark/BarrierTaskContext.scala
b/core/src/main/scala/org/apache/spark/BarrierTaskContext.scala
index 7f5d1f39d86c..49a5e77764a6 100644
--- a/core/src/main/scala/org/apache/spark/BarrierTaskContext.scala
+++ b/core/src/main/scala/org/apache/spark/BarrierTaskContext.scala
@@ -295,6 +295,15 @@ class BarrierTaskContext private[spark] (
: T = {
taskContext.createResourceUninterruptibly(resourceBuilder)
}
+
+ override def addPostStatusUpdateListener(listener:
PostStatusUpdateListener): TaskContext = {
+ taskContext.addPostStatusUpdateListener(listener)
+ this
+ }
+
+ override private[spark] def invokePostStatusUpdateListeners(): Unit = {
+ taskContext.invokePostStatusUpdateListeners()
+ }
}
@Experimental
diff --git a/core/src/main/scala/org/apache/spark/MapOutputTracker.scala
b/core/src/main/scala/org/apache/spark/MapOutputTracker.scala
index 41a1b51a4315..51ba146f728f 100644
--- a/core/src/main/scala/org/apache/spark/MapOutputTracker.scala
+++ b/core/src/main/scala/org/apache/spark/MapOutputTracker.scala
@@ -17,7 +17,7 @@
package org.apache.spark
-import java.io.{ByteArrayInputStream, InputStream, IOException,
ObjectInputStream, ObjectOutputStream}
+import java.io._
import java.nio.ByteBuffer
import java.util.concurrent.{ConcurrentHashMap, LinkedBlockingQueue,
ThreadPoolExecutor, TimeUnit}
import java.util.concurrent.locks.ReentrantReadWriteLock
@@ -105,6 +105,32 @@ private class ShuffleStatus(
*/
private[spark] val checksumMismatchIndices: Set[Int] = Set()
+ /**
+ * Set of stale pushed partition indexes for this shuffle. Each entry is a
partitionId (which
+ * equals mapIndex, not MapStatus.mapId). When task retry or speculation
causes multiple
+ * attempts for the same map output to push, the merger may include data
from a stale attempt.
+ * We record the stale partition indexes here so the reduce side can check
chunkBitmaps and
+ * fallback if stale data is present in a merged block.
+ */
+ private[this] val staleMapIndexes = new java.util.HashSet[Int]()
+
+ /**
+ * Mark a partition as having stale (redundant) push attempts. Called from
TaskSetManager when it
+ * detects that multiple task attempts for the same map output pushed data
to the merger.
+ * @param mapIndex the partition index (== mapIndex) of the stale
(redundant) attempt;
+ * this is NOT MapStatus.mapId
+ */
+ def markStalePushedPartition(mapIndex: Int): Unit = withWriteLock {
+ staleMapIndexes.add(mapIndex)
+ }
+
+ /**
+ * Get all stale pushed partition indexes for this shuffle. Returns empty
set if none exist.
+ */
+ def getStaleMapIndexes: Set[Int] = withReadLock {
+ new java.util.HashSet[Int](staleMapIndexes).asScala
+ }
+
/**
* MergeStatus for each shuffle partition when push-based shuffle is
enabled. The index of the
* array is the shuffle partition id (reduce id). Each value in the array is
the MergeStatus for
@@ -413,7 +439,7 @@ private class ShuffleStatus(
broadcastManager: BroadcastManager,
isLocal: Boolean,
minBroadcastSize: Int,
- conf: SparkConf): (Array[Byte], Array[Byte]) = {
+ conf: SparkConf): (Array[Byte], Array[Byte], Array[Byte]) = {
val mapStatusesBytes: Array[Byte] =
serializedMapStatus(broadcastManager, isLocal, minBroadcastSize, conf)
var mergeStatusesBytes: Array[Byte] = null
@@ -437,7 +463,12 @@ private class ShuffleStatus(
// `withWriteLock`.
mergeStatusesBytes = cachedSerializedMergeStatus
}
- (mapStatusesBytes, mergeStatusesBytes)
+
+ // Serialize staleMapIndexes set for reduce-side stale chunk-level
detection.
+ val staleMapIndexBytes = withReadLock {
+ MapOutputTracker.serializeStaleMapIndexes(staleMapIndexes)
+ }
+ (mapStatusesBytes, mergeStatusesBytes, staleMapIndexBytes)
}
// Used in testing.
@@ -688,6 +719,12 @@ private[spark] abstract class MapOutputTracker(conf:
SparkConf) extends Logging
*/
def getShufflePushMergerLocations(shuffleId: Int): Seq[BlockManagerId]
+ /**
+ * Get all stale pushed partition indexes for a given shuffle.
+ * Called from PushBasedFetchHelper on the reduce side for chunk-level
detection.
+ */
+ def getStaleMapIndexes(shuffleId: Int): Set[Int]
+
/**
* Deletes map output status information for the specified shuffle stage.
*/
@@ -958,6 +995,26 @@ private[spark] class MapOutputTrackerMaster(
incrementEpoch()
}
+ /**
+ * Mark a partition as having stale (redundant) push attempts and bump the
epoch so that
+ * reducers with a cached (empty) stale set are forced to re-fetch. Without
the epoch bump,
+ * a reducer that already fetched its merge statuses before this mark would
keep its stale
+ * set cached and never see the new mark, causing layer-3 fallback to
silently not fire.
+ *
+ * @param shuffleId the shuffle id.
+ * @param mapIndex the partition index (== mapIndex, not MapStatus.mapId) of
the stale
+ * (redundant) attempt; this is NOT MapStatus.mapId
+ */
+ def markStalePushedPartition(shuffleId: Int, mapIndex: Int): Unit = {
+ shuffleStatuses.get(shuffleId).foreach { shuffleStatus =>
+ shuffleStatus.markStalePushedPartition(mapIndex)
+ // Bump the epoch so reducers with a cached stale set are forced to
re-fetch.
+ // A reducer that fetched its merge statuses before this mark otherwise
keeps
+ // its (empty) stale set and layer-3 fallback would not fire.
+ incrementEpoch()
+ }
+ }
+
/** Check if the given shuffle is being tracked */
def containsShuffle(shuffleId: Int): Boolean =
shuffleStatuses.contains(shuffleId)
@@ -1260,6 +1317,10 @@ private[spark] class MapOutputTrackerMaster(
shuffleStatuses.get(shuffleId).map(_.getShufflePushMergerLocations).getOrElse(Seq.empty)
}
+ override def getStaleMapIndexes(shuffleId: Int): Set[Int] = {
+
shuffleStatuses.get(shuffleId).map(_.getStaleMapIndexes).getOrElse(Set.empty)
+ }
+
override def stop(): Unit = {
mapOutputTrackerMasterMessages.offer(PoisonPill)
threadpool.shutdown()
@@ -1305,6 +1366,28 @@ private[spark] class MapOutputTrackerWorker(conf:
SparkConf) extends MapOutputTr
*/
val shufflePushMergerLocations = new ConcurrentHashMap[Int,
Seq[BlockManagerId]]().asScala
+ /**
+ * Tracks stale pushed partition indexes per shuffle for speculation
stale-push detection.
+ * Populated during getStatuses fetch from driver, used by
PushBasedFetchHelper
+ * to check chunkBitmaps for stale data at chunk level.
+ */
+ private[spark] val staleMapIndexes =
+ new ConcurrentHashMap[Int, java.util.HashSet[Int]]().asScala
+
+ /**
+ * Get all stale pushed partition indexes for a given shuffle. Returns empty
set if none exist.
+ * Called from PushBasedFetchHelper on the reduce side for chunk-level
stale-data detection.
+ */
+ def getStaleMapIndexes(shuffleId: Int): Set[Int] = {
+ val dupSetOpt = staleMapIndexes.get(shuffleId)
+ dupSetOpt match {
+ case Some(dupSet) =>
+ val result = new java.util.HashSet[Int](dupSet) // defensive copy
+ result.asScala
+ case None => Set.empty[Int]
+ }
+ }
+
/**
* A [[KeyLock]] whose key is a shuffle id to ensure there is only one
thread fetching
* the same shuffle block.
@@ -1356,12 +1439,13 @@ private[spark] class MapOutputTrackerWorker(conf:
SparkConf) extends MapOutputTr
s"mappers $startMapIndex-$actualEndMapIndex, partitions
$startPartition-$endPartition")
MapOutputTracker.convertMapStatuses(
shuffleId, startPartition, endPartition, mapOutputStatuses,
startMapIndex,
- actualEndMapIndex, Option(mergedOutputStatuses))
+ actualEndMapIndex, Option(mergedOutputStatuses))
} catch {
case e: MetadataFetchFailedException =>
// We experienced a fetch failure so our mapStatuses cache is
outdated; clear it:
mapStatuses.clear()
mergeStatuses.clear()
+ staleMapIndexes.clear()
throw e
}
}
@@ -1386,6 +1470,7 @@ private[spark] class MapOutputTrackerWorker(conf:
SparkConf) extends MapOutputTr
case e: MetadataFetchFailedException =>
mapStatuses.clear()
mergeStatuses.clear()
+ staleMapIndexes.clear()
throw e
}
}
@@ -1407,6 +1492,7 @@ private[spark] class MapOutputTrackerWorker(conf:
SparkConf) extends MapOutputTr
case e: MetadataFetchFailedException =>
mapStatuses.clear()
mergeStatuses.clear()
+ staleMapIndexes.clear()
throw e
}
}
@@ -1456,13 +1542,20 @@ private[spark] class MapOutputTrackerWorker(conf:
SparkConf) extends MapOutputTr
if (fetchedMapStatuses == null || fetchedMergeStatuses == null) {
logInfo(log"Doing the fetch; tracker endpoint = " +
log"${MDC(RPC_ENDPOINT_REF, trackerEndpoint)}")
- val fetchedBytes =
- askTracker[(Array[Byte],
Array[Byte])](GetMapAndMergeResultStatuses(shuffleId))
+ val fetchedBytes = askTracker[(Array[Byte], Array[Byte],
Array[Byte])](
+ GetMapAndMergeResultStatuses(shuffleId))
try {
fetchedMapStatuses =
MapOutputTracker.deserializeOutputStatuses[MapStatus](fetchedBytes._1, conf)
fetchedMergeStatuses =
MapOutputTracker.deserializeOutputStatuses[MergeStatus](fetchedBytes._2, conf)
+ // Deserialize staleMapIndexes set for speculation stale-push
detection.
+ val deserializedStale =
MapOutputTracker.deserializeStaleMapIndexes(fetchedBytes._3)
+ staleMapIndexes.put(shuffleId, deserializedStale)
+ if (!deserializedStale.isEmpty) {
+ logInfo(s"Got stale pushed partition indexes for shuffle
$shuffleId: " +
+ deserializedStale.asScala.mkString(","))
+ }
} catch {
case e: SparkException =>
throw new MetadataFetchFailedException(shuffleId, -1,
@@ -1519,6 +1612,7 @@ private[spark] class MapOutputTrackerWorker(conf:
SparkConf) extends MapOutputTr
mapStatuses.remove(shuffleId)
mergeStatuses.remove(shuffleId)
shufflePushMergerLocations.remove(shuffleId)
+ staleMapIndexes.remove(shuffleId)
}
/**
@@ -1534,6 +1628,7 @@ private[spark] class MapOutputTrackerWorker(conf:
SparkConf) extends MapOutputTr
mapStatuses.clear()
mergeStatuses.clear()
shufflePushMergerLocations.clear()
+ staleMapIndexes.clear()
}
}
}
@@ -1644,6 +1739,48 @@ private[spark] object MapOutputTracker extends Logging {
}
}
+ /**
+ * Serialize a set of stale pushed partition indexes into a compact byte
array.
+ * Uses DataOutputStream for a simple, efficient binary format:
+ * [int: count][int: mapIndex1][int: mapIndex2]...
+ * This is intentionally lightweight compared to serializeOutputStatuses
because
+ * the set is typically small (only non-empty during speculation retries).
+ */
+ def serializeStaleMapIndexes(
+ mapIndexes: java.util.HashSet[Int]): Array[Byte] = {
+ val out = new ApacheByteArrayOutputStream()
+ val dataOut = new DataOutputStream(out)
+ Utils.tryWithSafeFinally {
+ dataOut.writeInt(mapIndexes.size())
+ val iter = mapIndexes.iterator()
+ while (iter.hasNext) {
+ dataOut.writeInt(iter.next())
+ }
+ } {
+ dataOut.close()
+ }
+ out.toByteArray
+ }
+
+ /**
+ * Deserialize a byte array produced by [[serializeStaleMapIndexes]] back
into a HashSet.
+ */
+ def deserializeStaleMapIndexes(bytes: Array[Byte]): java.util.HashSet[Int] =
{
+ val dupMapIdIn = new DataInputStream(new ByteArrayInputStream(bytes))
+ Utils.tryWithSafeFinally {
+ val numEntries = dupMapIdIn.readInt()
+ val dupSet = new java.util.HashSet[Int]()
+ var i = 0
+ while (i < numEntries) {
+ dupSet.add(dupMapIdIn.readInt())
+ i += 1
+ }
+ dupSet
+ } {
+ dupMapIdIn.close()
+ }
+ }
+
/**
* Given an array of map statuses and a range of map output partitions,
returns a sequence that,
* for each block manager ID, lists the shuffle block IDs and corresponding
shuffle block sizes
diff --git a/core/src/main/scala/org/apache/spark/TaskContext.scala
b/core/src/main/scala/org/apache/spark/TaskContext.scala
index 11f1e1e63f11..a4f68d76128d 100644
--- a/core/src/main/scala/org/apache/spark/TaskContext.scala
+++ b/core/src/main/scala/org/apache/spark/TaskContext.scala
@@ -27,7 +27,7 @@ import org.apache.spark.metrics.source.Source
import org.apache.spark.resource.ResourceInformation
import org.apache.spark.scheduler.Task
import org.apache.spark.shuffle.FetchFailedException
-import org.apache.spark.util.{AccumulatorV2, TaskCompletionListener,
TaskFailureListener, TaskInterruptListener}
+import org.apache.spark.util.{AccumulatorV2, PostStatusUpdateListener,
TaskCompletionListener, TaskFailureListener, TaskInterruptListener}
object TaskContext {
@@ -195,6 +195,23 @@ abstract class TaskContext extends Serializable {
})
}
+ /**
+ * Adds a listener to be invoked after the task's status update has been
sent to the driver.
+ * This is useful for operations that should only begin after the driver has
been notified
+ * of the task's result. For example, push-based shuffle block push can use
this to
+ * ensure the driver processes the task result before any push data reaches
the merger,
+ * avoiding stale data being merged without detection.
+ *
+ * The callback runs on the same executor thread that sends the status
update.
+ */
+ private[spark] def addPostStatusUpdateListener(listener:
PostStatusUpdateListener): TaskContext
+
+ /**
+ * Invokes all registered post-status-update listeners. Called by Executor
after sending
+ * the task's status update to the driver.
+ */
+ private[spark] def invokePostStatusUpdateListeners(): Unit
+
/** Runs a task with this context, ensuring failure and completion listeners
get triggered. */
private[spark] def runTaskWithListeners[T](task: Task[T]): T = {
try {
diff --git a/core/src/main/scala/org/apache/spark/TaskContextImpl.scala
b/core/src/main/scala/org/apache/spark/TaskContextImpl.scala
index 24202d92601d..9bf9ce0d9de2 100644
--- a/core/src/main/scala/org/apache/spark/TaskContextImpl.scala
+++ b/core/src/main/scala/org/apache/spark/TaskContextImpl.scala
@@ -76,6 +76,13 @@ private[spark] class TaskContextImpl(
/** List of callback functions to execute when the task is interrupted. */
@transient private val onInterruptCallbacks = new
Stack[TaskInterruptListener]
+ /**
+ * List of callback functions to invoke after the task's status update has
been sent
+ * to the driver. Used for operations like push-based shuffle block push
that should only
+ * begin after the driver has been notified of the task result.
+ */
+ @transient private val onPostStatusUpdateCallbacks = new
Stack[PostStatusUpdateListener]
+
/**
* The thread currently executing task completion or failure listeners, if
any.
*
@@ -144,6 +151,19 @@ private[spark] class TaskContextImpl(
this
}
+ override def addPostStatusUpdateListener(listener:
PostStatusUpdateListener): this.type = {
+ synchronized {
+ onPostStatusUpdateCallbacks.push(listener)
+ }
+ this
+ }
+
+ private[spark] def invokePostStatusUpdateListeners(): Unit = {
+ invokeListeners(onPostStatusUpdateCallbacks, "PostStatusUpdateListener",
None) {
+ _.onStatusUpdateSent(this)
+ }
+ }
+
override def resourcesJMap(): java.util.Map[String, ResourceInformation] = {
resources.asJava
}
diff --git a/core/src/main/scala/org/apache/spark/executor/Executor.scala
b/core/src/main/scala/org/apache/spark/executor/Executor.scala
index 805ec68a8947..1790c781ecfc 100644
--- a/core/src/main/scala/org/apache/spark/executor/Executor.scala
+++ b/core/src/main/scala/org/apache/spark/executor/Executor.scala
@@ -1027,6 +1027,9 @@ private[spark] class Executor(
setTaskFinishedAndClearInterruptStatus()
plugins.foreach(_.onTaskSucceeded())
execBackend.statusUpdate(taskId, TaskState.FINISHED, serializedResult)
+ Utils.tryLogNonFatalError {
+ task.context.invokePostStatusUpdateListeners()
+ }
} catch {
case t: TaskKilledException =>
logInfo(log"Executor killed ${MDC(TASK_NAME, taskName)}," +
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 5c077a7a3bbb..2e8e6628ea5d 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
@@ -824,6 +824,8 @@ private[spark] class TaskSetManager(
val index = info.index
// Check if any other attempt succeeded before this and this attempt has
not been handled
if (successful(index) && killedByOtherAttempt.contains(tid)) {
+ // For ShuffleMapTasks, detect speculation duplicate before handling as
killed.
+ detectStalePushIfShuffleTask(tid, info.index, result)
// Undo the effect on calculatedTasks and totalResultSize made earlier
when
// checking if can fetch more results
calculatedTasks -= 1
@@ -930,6 +932,34 @@ private[spark] class TaskSetManager(
taskInfoWithAccumulables)
}
+ /**
+ * For ShuffleMapTasks, detect stale push: if a partition already has
+ * a registered MapStatus with a different mapId, it means another attempt
for the same
+ * partition also pushed data to the merger. Mark this partition so that
reducers will
+ * skip the merged block and fallback to unmerged blocks.
+ *
+ * This is called from handleSuccessfulTask for late-arriving or killed
attempt results,
+ * where the task result won't be forwarded to DAGScheduler (so
DAGScheduler's own
+ * stale detection won't cover these cases).
+ */
+ private def detectStalePushIfShuffleTask(
+ tid: Long, index: Int, result: DirectTaskResult[_]): Unit = {
+ if (!isShuffleMapTasks || shuffleId.isEmpty) {
+ return
+ }
+ // This method is only called for late-arriving or killed attempts,
meaning the partition
+ // already has a successful attempt registered. Any MapStatus arriving
here is from a
+ // stale (redundant) attempt that also pushed data to the merger.
+ val mapStatus = result.value().asInstanceOf[MapStatus]
+ val sid = shuffleId.get
+ val partitionId = tasks(index).partitionId
+ // Mark its mapIndex (== partitionId, NOT MapStatus.mapId) as stale so
reducers can detect
+ // the stale data in merged block chunks via chunkBitmaps and fallback if
needed.
+ logInfo(s"[StalePush] Late/killed attempt tid=$tid for
partition=$partitionId " +
+ s"(index=$index), attemptMapId=${mapStatus.mapId}. Marked as stale
push.")
+ sched.mapOutputTracker.markStalePushedPartition(sid, partitionId)
+ }
+
private[scheduler] def markPartitionCompleted(partitionId: Int): Unit = {
partitionToIndex.get(partitionId).foreach { index =>
if (!successful(index)) {
diff --git
a/core/src/main/scala/org/apache/spark/shuffle/ShuffleWriteProcessor.scala
b/core/src/main/scala/org/apache/spark/shuffle/ShuffleWriteProcessor.scala
index 47d54ae4f10b..c397ac03ed92 100644
--- a/core/src/main/scala/org/apache/spark/shuffle/ShuffleWriteProcessor.scala
+++ b/core/src/main/scala/org/apache/spark/shuffle/ShuffleWriteProcessor.scala
@@ -21,6 +21,7 @@ import org.apache.spark.{ShuffleDependency, SparkEnv,
TaskContext}
import org.apache.spark.internal.Logging
import org.apache.spark.internal.LogKeys.{NUM_MERGER_LOCATIONS, SHUFFLE_ID,
STAGE_ID}
import org.apache.spark.scheduler.MapStatus
+import org.apache.spark.util.PostStatusUpdateListener
/**
* The interface for customizing shuffle write process. The driver create a
ShuffleWriteProcessor
@@ -79,8 +80,29 @@ private[spark] class ShuffleWriteProcessor extends
Serializable with Logging {
log" with shuffle ID ${MDC(SHUFFLE_ID, dep.shuffleId)}")
logDebug(s"Starting pushing blocks for the task
${context.taskAttemptId()}")
val dataFile = resolver.getDataFile(dep.shuffleId, mapId)
- new ShuffleBlockPusher(SparkEnv.get.conf)
- .initiateBlockPush(dataFile, writer.getPartitionLengths(),
dep, mapIndex)
+ val blockPusher = new ShuffleBlockPusher(SparkEnv.get.conf)
+ // Register a post-status-update listener to defer push until
after the task
+ // result has been sent to the driver. This ensures the driver
processes the
+ // task result (and can mark stale partitions from speculative
duplicates)
+ // before any push data reaches the merger, avoiding stale data
being merged
+ // without detection.
+ //
+ // The listener callback runs on the Task thread and only does a
lightweight
+ // submitTask; actual push I/O runs on BLOCK_PUSHER_POOL threads.
+ context.addPostStatusUpdateListener(new PostStatusUpdateListener
{
+ override def onStatusUpdateSent(context: TaskContext): Unit = {
+ if (!context.isInterrupted() && !context.isFailed()) {
+ logDebug(s"Task ${context.taskAttemptId()} status update
sent, " +
+ s"proceeding with shuffle block push for shuffle
${dep.shuffleId}")
+ blockPusher.initiateBlockPush(
+ dataFile, writer.getPartitionLengths(), dep, mapIndex)
+ } else {
+ logInfo(s"Task ${context.taskAttemptId()} was " +
+ s"${if (context.isInterrupted()) "killed" else
"failed"}, " +
+ s"skipping shuffle block push for shuffle
${dep.shuffleId}")
+ }
+ }
+ })
case _ =>
}
}
diff --git
a/core/src/main/scala/org/apache/spark/storage/PushBasedFetchHelper.scala
b/core/src/main/scala/org/apache/spark/storage/PushBasedFetchHelper.scala
index 16a68bbf50e2..aa0413be38a1 100644
--- a/core/src/main/scala/org/apache/spark/storage/PushBasedFetchHelper.scala
+++ b/core/src/main/scala/org/apache/spark/storage/PushBasedFetchHelper.scala
@@ -19,7 +19,6 @@ package org.apache.spark.storage
import java.util.concurrent.TimeUnit
-import scala.collection
import scala.collection.mutable
import scala.collection.mutable.ArrayBuffer
import scala.util.{Failure, Success}
@@ -166,9 +165,16 @@ private class PushBasedFetchHelper(
meta: MergedBlockMeta): Unit = {
logDebug(s"Received the meta of push-merged block for ($shuffleId,
$shuffleMergeId," +
s" $reduceId) from ${req.address.host}:${req.address.port}")
+ val mergedBlock = ShuffleMergedBlockId(shuffleId, shuffleMergeId,
reduceId)
try {
-
iterator.addToResultsQueue(PushMergedRemoteMetaFetchResult(shuffleId,
shuffleMergeId,
- reduceId, sizeMap((shuffleId, reduceId)), meta.readChunkBitmaps(),
address))
+ val chunkBitmaps = meta.readChunkBitmaps()
+ if (checkStaleMapIdInMergedBlock(mergedBlock, address,
chunkBitmaps)) {
+
iterator.addToResultsQueue(PushMergedRemoteMetaFetchResult(shuffleId,
shuffleMergeId,
+ reduceId, sizeMap((shuffleId, reduceId)), chunkBitmaps, address))
+ } else {
+
iterator.addToResultsQueue(PushMergedRemoteMetaFailedFetchResult(shuffleId,
+ shuffleMergeId, reduceId, address))
+ }
} catch {
case exception: Exception =>
logError(log"Failed to parse the meta of push-merged block for (" +
@@ -273,9 +279,15 @@ private class PushBasedFetchHelper(
try {
val shuffleBlockId = blockId.asInstanceOf[ShuffleMergedBlockId]
val chunksMeta = blockManager.getLocalMergedBlockMeta(shuffleBlockId,
localDirs)
- iterator.addToResultsQueue(PushMergedLocalMetaFetchResult(
- shuffleBlockId.shuffleId, shuffleBlockId.shuffleMergeId,
- shuffleBlockId.reduceId, chunksMeta.readChunkBitmaps(), localDirs))
+ val chunkBitmaps = chunksMeta.readChunkBitmaps()
+ if (checkStaleMapIdInMergedBlock(shuffleBlockId, blockManagerId,
chunkBitmaps)) {
+ iterator.addToResultsQueue(PushMergedLocalMetaFetchResult(
+ shuffleBlockId.shuffleId, shuffleBlockId.shuffleMergeId,
+ shuffleBlockId.reduceId, chunkBitmaps, localDirs))
+ } else {
+ iterator.addToResultsQueue(FallbackOnPushMergedFailureResult(
+ blockId, blockManagerId, 0, isNetworkReqDone = false))
+ }
} catch {
case e: Exception =>
// If we see an exception with reading a push-merged-local meta, we
fallback to
@@ -288,6 +300,36 @@ private class PushBasedFetchHelper(
}
}
+ /**
+ * Check whether a push-merged block contains data from stale (duplicate)
task attempts.
+ * When speculation is enabled, multiple attempts for the same map output
may both push data
+ * to the merger. The merger may include data from both attempts in the same
merged block,
+ * but the driver only tracks one as the canonical MapStatus. We detect this
by checking
+ * if any stale pushed map index appears in the server-side chunkBitmaps.
+ *
+ * @param shuffleBlockId ShuffleMergedBlockId to be checked
+ * @param address BlockManagerId of push-based shuffle service
+ * @param chunkBitmaps Chunks bitmap from push-based shuffle service side
+ * @return false if any stale-marked mapIndex is present in this block
(forcing fallback),
+ * true otherwise
+ */
+ private[this] def checkStaleMapIdInMergedBlock(
+ shuffleBlockId: ShuffleMergedBlockId,
+ address: BlockManagerId,
+ chunkBitmaps: Array[RoaringBitmap]): Boolean = {
+ val staleMapIndexes =
+ mapOutputTracker.getStaleMapIndexes(shuffleBlockId.shuffleId)
+ if (staleMapIndexes.isEmpty) return true
+ val mergedBlockBitmap = new RoaringBitmap()
+ chunkBitmaps.foreach(mergedBlockBitmap.or)
+ val hasStale = staleMapIndexes.exists(id => mergedBlockBitmap.contains(id))
+ if (hasStale) {
+ logWarning(s"Found stale pushed map indexes in merged block
$shuffleBlockId from" +
+ s" ${address.host}:${address.port}, falling back to fetch the original
blocks")
+ }
+ !hasStale
+ }
+
/**
* This is executed by the task thread when the `iterator.next()` is invoked
and the iterator
* processes a response of type:
diff --git a/core/src/main/scala/org/apache/spark/util/taskListeners.scala
b/core/src/main/scala/org/apache/spark/util/taskListeners.scala
index 39b7dbe66f23..b40d1a5fc1b5 100644
--- a/core/src/main/scala/org/apache/spark/util/taskListeners.scala
+++ b/core/src/main/scala/org/apache/spark/util/taskListeners.scala
@@ -54,6 +54,23 @@ trait TaskInterruptListener extends EventListener {
def onTaskInterrupted(context: TaskContext, reason: String): Unit
}
+/**
+ * Listener providing a callback to be invoked after the task's status update
has been sent to
+ * the driver. This is useful for operations that should only begin after the
driver has been
+ * notified of the task's result, such as push-based shuffle block push, where
starting push
+ * before the driver processes the result can lead to stale data being merged
without detection.
+ *
+ * The callback runs on the same executor thread that sends the status update.
It should perform
+ * only lightweight work (e.g., submitting async I/O to another thread pool)
to avoid blocking
+ * the executor task thread.
+ *
+ * Listeners are guaranteed to execute sequentially in reverse order of
registration.
+ */
+private[spark]
+trait PostStatusUpdateListener extends EventListener {
+ def onStatusUpdateSent(context: TaskContext): Unit
+}
+
/**
* Exception thrown when there is an exception in executing the callback in
TaskCompletionListener.
diff --git a/core/src/test/scala/org/apache/spark/MapOutputTrackerSuite.scala
b/core/src/test/scala/org/apache/spark/MapOutputTrackerSuite.scala
index d2344b4e7291..2015ceb8b090 100644
--- a/core/src/test/scala/org/apache/spark/MapOutputTrackerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/MapOutputTrackerSuite.scala
@@ -637,7 +637,7 @@ class MapOutputTrackerSuite extends SparkFunSuite with
LocalSparkContext {
mapWorkerRpcEnv.setupEndpointRef(rpcEnv.address,
MapOutputTracker.ENDPOINT_NAME)
val fetchedBytes = mapWorkerTracker.trackerEndpoint
- .askSync[(Array[Byte], Array[Byte])](GetMapAndMergeResultStatuses(20))
+ .askSync[(Array[Byte], Array[Byte],
Array[Byte])](GetMapAndMergeResultStatuses(20))
assert(masterTracker.getNumAvailableMergeResults(20) == 1)
assert(masterTracker.getNumAvailableOutputs(20) == 100)
@@ -1166,4 +1166,56 @@ class MapOutputTrackerSuite extends SparkFunSuite with
LocalSparkContext {
rpcEnv.shutdown()
}
}
+
+ test("SPARK-57491: staleMapIndexes propagated from driver to worker via
getStatuses") {
+ val rpcEnv = createRpcEnv("test")
+ val masterTracker = newTrackerMaster()
+ try {
+ masterTracker.trackerEndpoint =
rpcEnv.setupEndpoint(MapOutputTracker.ENDPOINT_NAME,
+ new MapOutputTrackerMasterEndpoint(rpcEnv, masterTracker, conf))
+ // Simulate driver side: register shuffle
+ masterTracker.registerShuffle(0, 4, 2)
+ masterTracker.registerMapOutput(0, 0,
+ MapStatus(BlockManagerId("exec-1", "hostA", 1000), Array(100L, 200L),
0))
+ masterTracker.registerMapOutput(0, 1,
+ MapStatus(BlockManagerId("exec-2", "hostB", 1000), Array(300L, 400L),
1))
+
+ val initialEpoch = masterTracker.getEpoch
+
+ // Mark partition 1 as stale via MapOutputTrackerMaster: marks the
partition and bumps
+ // the epoch so reducers with a cached (empty) stale set are forced to
re-fetch
+ masterTracker.markStalePushedPartition(0, 1)
+
+ // Verify epoch was bumped so reducer-side stale set cache is invalidated
+ assert(masterTracker.getEpoch > initialEpoch,
+ s"Expected epoch to be bumped past $initialEpoch, got
${masterTracker.getEpoch}")
+
+ val shuffleStatus = masterTracker.shuffleStatuses(0)
+
+ // Verify the stale mark was recorded on the shuffle status
+ assert(shuffleStatus.getStaleMapIndexes.contains(1))
+
+ // Serialize statuses including staleMapIndexes (simulates what driver
sends to executors)
+ val (mapBytes, mergeBytes, staleBytes) =
+ shuffleStatus.serializedMapAndMergeStatus(
+ masterTracker.broadcastManager, isLocal = true, minBroadcastSize =
Int.MaxValue, conf)
+
+ // Simulate worker side: deserialize and verify staleMapIndexes are
received
+ val deserializedStale =
MapOutputTracker.deserializeStaleMapIndexes(staleBytes)
+ assert(deserializedStale.contains(1))
+
+ // Also verify via MapOutputTrackerWorker (the path used by real
executors)
+ val workerTracker = new MapOutputTrackerWorker(conf)
+ // Manually populate what the worker would get from a getStatuses fetch
+ workerTracker.staleMapIndexes.put(0, {
+ val s = new JHashSet[Int]()
+ s.addAll(deserializedStale)
+ s
+ })
+ assert(workerTracker.getStaleMapIndexes(0).contains(1))
+ } finally {
+ masterTracker.stop()
+ rpcEnv.shutdown()
+ }
+ }
}
diff --git
a/core/src/test/scala/org/apache/spark/scheduler/TaskContextSuite.scala
b/core/src/test/scala/org/apache/spark/scheduler/TaskContextSuite.scala
index ac8d000d8a44..d9bfbe2b5f24 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/TaskContextSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/TaskContextSuite.scala
@@ -768,6 +768,76 @@ class TaskContextSuite extends SparkFunSuite with
BeforeAndAfter with LocalSpark
assert(isFailed)
}
+ test("SPARK-57491: invokePostStatusUpdateListeners triggers
PostStatusUpdateListener") {
+ // Basic invocation - listener is called when
invokePostStatusUpdateListeners is invoked
+ val context1 = TaskContext.empty()
+ var invoked = false
+ context1.addPostStatusUpdateListener(new PostStatusUpdateListener {
+ override def onStatusUpdateSent(context: TaskContext): Unit = invoked =
true
+ })
+ assert(!invoked)
+ context1.invokePostStatusUpdateListeners()
+ assert(invoked)
+
+ // Multiple listeners are called in reverse registration order
+ val context2 = TaskContext.empty()
+ val callOrder = ArrayBuffer.empty[String]
+ context2.addPostStatusUpdateListener(new PostStatusUpdateListener {
+ override def onStatusUpdateSent(context: TaskContext): Unit = callOrder
+= "first"
+ })
+ context2.addPostStatusUpdateListener(new PostStatusUpdateListener {
+ override def onStatusUpdateSent(context: TaskContext): Unit = callOrder
+= "second"
+ })
+ context2.invokePostStatusUpdateListeners()
+ assert(callOrder === Seq("second", "first"))
+
+ // Listener exception is wrapped in TaskCompletionListenerException
+ val context3 = TaskContext.empty()
+ context3.addPostStatusUpdateListener(new PostStatusUpdateListener {
+ override def onStatusUpdateSent(context: TaskContext): Unit =
+ throw new RuntimeException("post-status-update failure")
+ })
+ val e = intercept[TaskCompletionListenerException] {
+ context3.invokePostStatusUpdateListeners()
+ }
+ assert(e.getMessage.contains("post-status-update failure"))
+
+ // All listeners are called even if one throws (same behavior as
completion listeners)
+ val context4 = TaskContext.empty()
+ val listener = mock(classOf[PostStatusUpdateListener])
+ context4.addPostStatusUpdateListener(new PostStatusUpdateListener {
+ override def onStatusUpdateSent(context: TaskContext): Unit =
+ throw new RuntimeException("bad listener")
+ })
+ context4.addPostStatusUpdateListener(listener)
+ context4.addPostStatusUpdateListener(new PostStatusUpdateListener {
+ override def onStatusUpdateSent(context: TaskContext): Unit =
+ throw new RuntimeException("another bad listener")
+ })
+
+ intercept[TaskCompletionListenerException] {
+ context4.invokePostStatusUpdateListeners()
+ }
+
+ verify(listener, times(1)).onStatusUpdateSent(any())
+
+ // invokePostStatusUpdateListeners is independent from
TaskCompletionListener
+ // (completion listener should NOT be triggered by
invokePostStatusUpdateListeners)
+ val context5 = TaskContext.empty()
+ var postStatusInvoked = false
+ var completionInvoked = false
+ context5.addPostStatusUpdateListener(new PostStatusUpdateListener {
+ override def onStatusUpdateSent(context: TaskContext): Unit =
postStatusInvoked = true
+ })
+ context5.addTaskCompletionListener(new TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit =
completionInvoked = true
+ })
+
+ context5.invokePostStatusUpdateListeners()
+ assert(postStatusInvoked)
+ assert(!completionInvoked) // completion listener should NOT be triggered
+ }
+
test("SPARK-46947: ensure the correct block manager is used to unroll memory
for task") {
import BlockManagerValidationPlugin._
BlockManagerValidationPlugin.resetState()
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 4fd710712221..299bce879ba2 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala
@@ -2921,6 +2921,89 @@ class TaskSetManagerSuite
s"\nCaptured logs:\n${logs.mkString("\n")}")
}
+ test("SPARK-57491: late-arriving speculative ShuffleMapTask marks stale
partitionId") {
+ sc = new SparkContext("local", "test")
+ sched = new FakeTaskScheduler(sc, ("exec1", "host1"), ("exec2", "host2"),
("exec3", "host3"))
+ sc.conf.set(config.SPECULATION_MULTIPLIER, 0.0)
+ sc.conf.set(config.SPECULATION_ENABLED, true)
+
+ val taskSet = FakeTask.createShuffleMapTaskSet(2, 0, 0,
+ Seq(TaskLocation("host1", "exec1")),
+ Seq(TaskLocation("host2", "exec2")))
+ val clock = new ManualClock()
+ val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES, clock
= clock)
+ val accumUpdatesByTask: Array[Seq[AccumulatorV2[_, _]]] =
taskSet.tasks.map { task =>
+ task.metrics.internalAccums
+ }
+
+ // Register shuffle in MapOutputTrackerMaster so
detectStalePushIfShuffleTask can find it
+ val mapOutputTrackerMaster = sched.mapOutputTracker
+ val shuffleId = taskSet.shuffleId.get
+ mapOutputTrackerMaster.registerShuffle(shuffleId, 2, 2)
+
+ // Offer resources for 2 tasks to start
+ val task0 = manager.resourceOffer("exec1", "host1", PROCESS_LOCAL)._1.get
+ val task1 = manager.resourceOffer("exec2", "host2", PROCESS_LOCAL)._1.get
+ assert(task0.index === 0)
+ assert(task1.index === 1)
+
+ // Advance clock so tasks have been running long enough (markFinished
requires time > 0)
+ clock.advance(1)
+
+ // Complete task 1 (partition 1) successfully with a MapStatus
+ val mapStatus1 = MapStatus(BlockManagerId("exec2", "host2", 2000),
Array(2L, 2L), mapTaskId = 1)
+ val result1 = createMapStatusTaskResult(mapStatus1, accumUpdatesByTask(1))
+ manager.handleSuccessfulTask(task1.taskId, result1)
+ assert(sched.endedTasks(task1.index) === Success)
+
+ // Advance clock so task 0 has been running long enough for speculation.
+ // checkSpeculatableTasks requires tasks to have been running > 0ms when
threshold is 0.
+ clock.advance(1)
+ assert(manager.checkSpeculatableTasks(0))
+ assert(sched.speculativeTasks.toSet === Set(0))
+
+ // Offer resource to start the speculative attempt for partition 0 on a
different host
+ val specTaskOption = manager.resourceOffer("exec3", "host3", ANY)._1
+ assert(specTaskOption.isDefined, "Expected speculative task to be
launched")
+ val specTask = specTaskOption.get
+ assert(specTask.index === 0)
+ assert(specTask.attemptNumber === 1)
+
+ // Replace backend with mock before completing original task 0, to handle
killTask call
+ sched.backend = mock(classOf[SchedulerBackend])
+ sched.dagScheduler.stop()
+ sched.dagScheduler = mock(classOf[DAGScheduler])
+
+ // Complete original task 0 (partition 0) - this will kill the speculative
attempt
+ val mapStatus0 = MapStatus(BlockManagerId("exec1", "host1", 1000),
Array(1L, 1L), mapTaskId = 0)
+ val result0 = createMapStatusTaskResult(mapStatus0, accumUpdatesByTask(0))
+ manager.handleSuccessfulTask(task0.taskId, result0)
+
+ // Verify no stale pushed map indexes yet (stale is only marked when late
result arrives)
+ assert(mapOutputTrackerMaster.getStaleMapIndexes(shuffleId).isEmpty)
+
+ // Now the speculative attempt's result arrives late. Since task 0 already
succeeded,
+ // handleSuccessfulTask will see successful(0)=true and
killedByOtherAttempt contains
+ // the speculative tid, triggering detectStalePushIfShuffleTask.
+ val specMapStatus = MapStatus(
+ BlockManagerId("exec3", "host3", 3000), Array(3L, 3L), mapTaskId = 999)
+ val specResult = createMapStatusTaskResult(specMapStatus,
accumUpdatesByTask(0))
+ manager.handleSuccessfulTask(specTask.taskId, specResult)
+
+ // Verify that partition 0 is now tracked as stale
+ val staleMapIndexes = mapOutputTrackerMaster.getStaleMapIndexes(shuffleId)
+ assert(staleMapIndexes.contains(0),
+ s"Expected staleMapIndexes to contain mapIndex 0, got $staleMapIndexes")
+ }
+
+ private def createMapStatusTaskResult(
+ mapStatus: MapStatus,
+ accumUpdates: Seq[AccumulatorV2[_, _]],
+ metricPeaks: Array[Long] = Array.empty): DirectTaskResult[MapStatus] = {
+ val valueSer = SparkEnv.get.serializer.newInstance()
+ new DirectTaskResult[MapStatus](valueSer.serialize(mapStatus),
accumUpdates, metricPeaks)
+ }
+
}
class FakeLongTasks(stageId: Int, partitionId: Int) extends FakeTask(stageId,
partitionId) {
diff --git
a/core/src/test/scala/org/apache/spark/storage/ShuffleBlockFetcherIteratorSuite.scala
b/core/src/test/scala/org/apache/spark/storage/ShuffleBlockFetcherIteratorSuite.scala
index 211de2e8729e..f9d2fdf941d0 100644
---
a/core/src/test/scala/org/apache/spark/storage/ShuffleBlockFetcherIteratorSuite.scala
+++
b/core/src/test/scala/org/apache/spark/storage/ShuffleBlockFetcherIteratorSuite.scala
@@ -59,6 +59,7 @@ class ShuffleBlockFetcherIteratorSuite extends SparkFunSuite {
mapOutputTracker = mock(classOf[MapOutputTracker])
when(mapOutputTracker.getMapSizesForMergeResult(any(), any(), any()))
.thenReturn(Seq.empty.iterator)
+
when(mapOutputTracker.getStaleMapIndexes(any())).thenReturn(mutable.Set.empty[Int])
}
private def doReturn(value: Any) = org.mockito.Mockito.doReturn(value,
Seq.empty: _*)
@@ -2076,4 +2077,84 @@ class ShuffleBlockFetcherIteratorSuite extends
SparkFunSuite {
assert(iterator.next()._1 === ShuffleBlockId(0, 1, 0))
assert(!iterator.hasNext)
}
+
+ test("SPARK-57491: fallback behavior when stale mapIndex is in/not in merged
block") {
+ val blockManager = createMockBlockManager()
+ val localHost = "test-local-host"
+ val localBmId = BlockManagerId("test-client", localHost, 1)
+ val pushMergedBmId = BlockManagerId(SHUFFLE_MERGER_IDENTIFIER, localHost,
1)
+ val localDirs = Array("test-merged-dir-1", "test-merged-dir-2")
+ initHostLocalDirManager(blockManager, Map(SHUFFLE_MERGER_IDENTIFIER ->
localDirs))
+
+ when(mapOutputTracker.getStaleMapIndexes(0)).thenReturn(mutable.Set(1))
+ when(mapOutputTracker.getMapSizesForMergeResult(0, 2))
+ .thenReturn(Seq((localBmId,
+ toBlockList(Seq(ShuffleBlockId(0, 0, 2), ShuffleBlockId(0, 1, 2)), 1L,
1))).iterator)
+ doReturn(createMockManagedBuffer(100)).when(blockManager)
+ .getLocalBlockData(ShuffleBlockId(0, 0, 2))
+ doReturn(createMockManagedBuffer(100)).when(blockManager)
+ .getLocalBlockData(ShuffleBlockId(0, 1, 2))
+
+ // Stale mapIndex IS in the merged block's bitmap -> should trigger
fallback
+ val bitmaps1 = Array(new RoaringBitmap)
+ bitmaps1(0).add(0) // chunk 0 has mapIndex 0
+ bitmaps1(0).add(1) // chunk 0 also has mapIndex 1 (stale)
+
+ val pushMergedBlockMeta1 = createMockPushMergedBlockMeta(bitmaps1.length,
bitmaps1)
+ when(blockManager.getLocalMergedBlockMeta(ShuffleMergedBlockId(0, 0, 2),
localDirs))
+ .thenReturn(pushMergedBlockMeta1)
+ doReturn(Seq(createMockManagedBuffer(100), createMockManagedBuffer(100)))
+ .when(blockManager).getLocalMergedBlockData(ShuffleMergedBlockId(0, 0,
2), localDirs)
+ doReturn(createMockManagedBuffer(100)).when(blockManager)
+ .getLocalBlockData(ShuffleBlockId(0, 3, 2))
+
+ val blocksByAddress1 = Map[BlockManagerId, Seq[(BlockId, Long, Int)]](
+ (localBmId, toBlockList(Seq(ShuffleBlockId(0, 3, 2)), 1L, 1)),
+ (pushMergedBmId, toBlockList(Seq(ShuffleMergedBlockId(0, 0, 2)), 200L,
+ SHUFFLE_PUSH_MAP_ID)))
+
+ val taskContext1 = TaskContext.empty()
+ val shuffleMetrics1 =
taskContext1.taskMetrics.createTempShuffleReadMetrics()
+ val iterator1 = createShuffleBlockIteratorWithDefaults(blocksByAddress1,
+ blockManager = Some(blockManager),
+ shuffleMetrics = Some(shuffleMetrics1))
+
+ val (id1, _) = iterator1.next()
+ assert(id1 === ShuffleBlockId(0, 3, 2))
+
+ val (id2, _) = iterator1.next()
+ assert(id2 === ShuffleBlockId(0, 0, 2))
+ val (id3, _) = iterator1.next()
+ assert(id3 === ShuffleBlockId(0, 1, 2))
+ assert(!iterator1.hasNext)
+ assert(shuffleMetrics1.mergedFetchFallbackCount === 1)
+
+ // Stale mapIndex is NOT in the merged block's bitmap -> should NOT
trigger fallback
+ val bitmaps2 = Array(new RoaringBitmap, new RoaringBitmap)
+ bitmaps2(0).add(0) // chunk 0 has mapIndex 0
+ bitmaps2(1).add(2) // chunk 1 has mapIndex 2
+
+ val pushMergedBlockMeta2 = createMockPushMergedBlockMeta(bitmaps2.length,
bitmaps2)
+ when(blockManager.getLocalMergedBlockMeta(ShuffleMergedBlockId(0, 0, 2),
localDirs))
+ .thenReturn(pushMergedBlockMeta2)
+ doReturn(Seq(createMockManagedBuffer(100), createMockManagedBuffer(100)))
+ .when(blockManager).getLocalMergedBlockData(ShuffleMergedBlockId(0, 0,
2), localDirs)
+
+ val blocksByAddress2 = Map[BlockManagerId, Seq[(BlockId, Long, Int)]](
+ (pushMergedBmId, toBlockList(Seq(ShuffleMergedBlockId(0, 0, 2)), 200L,
+ SHUFFLE_PUSH_MAP_ID)))
+
+ val taskContext2 = TaskContext.empty()
+ val shuffleMetrics2 =
taskContext2.taskMetrics.createTempShuffleReadMetrics()
+ val iterator2 = createShuffleBlockIteratorWithDefaults(blocksByAddress2,
+ blockManager = Some(blockManager),
+ shuffleMetrics = Some(shuffleMetrics2))
+
+ val (id4, _) = iterator2.next()
+ assert(id4.isInstanceOf[ShuffleBlockChunkId])
+ val (id5, _) = iterator2.next()
+ assert(id5.isInstanceOf[ShuffleBlockChunkId])
+ assert(!iterator2.hasNext)
+ assert(shuffleMetrics2.mergedFetchFallbackCount === 0)
+ }
}
diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala
index 48ba1d418b65..643e4dca984e 100644
--- a/project/MimaExcludes.scala
+++ b/project/MimaExcludes.scala
@@ -83,7 +83,10 @@ object MimaExcludes {
// methods of V2ExpressionSQLBuilder. MySQLDialect is private, so this is
not a public API.
ProblemFilters.exclude[InaccessibleMethodProblem]("org.apache.spark.sql.connector.util.V2ExpressionSQLBuilder.visitStartsWith"),
ProblemFilters.exclude[InaccessibleMethodProblem]("org.apache.spark.sql.connector.util.V2ExpressionSQLBuilder.visitEndsWith"),
-
ProblemFilters.exclude[InaccessibleMethodProblem]("org.apache.spark.sql.connector.util.V2ExpressionSQLBuilder.visitContains")
+
ProblemFilters.exclude[InaccessibleMethodProblem]("org.apache.spark.sql.connector.util.V2ExpressionSQLBuilder.visitContains"),
+ // [SPARK-57491][CORE] Add PostStatusUpdateListener to TaskContext for
stale push detection
+
ProblemFilters.exclude[ReversedMissingMethodProblem]("org.apache.spark.TaskContext.addPostStatusUpdateListener"),
+
ProblemFilters.exclude[ReversedMissingMethodProblem]("org.apache.spark.TaskContext.invokePostStatusUpdateListeners")
)
// Exclude rules for 4.1.x from 4.0.0
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
