From e6894190fb0662bf1b63a3645f6f3700eb828d2b Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sat, 11 Jul 2026 05:33:58 +0000 Subject: [PATCH 01/39] [SPARK-XXXXX][CORE] Fail a pipelined group atomically on any member task failure A pipelined-group member reads/writes a transient shuffle that cannot be re-read in isolation, so a single member task failure must fail the whole group -- which fails the job, letting the caller (the streaming batch loop) rerun -- rather than be retried per-task. TaskSet gains an isPipelined flag, set by DAGScheduler.isPipelinedGroupMember (true if the stage produces a PipelinedShuffleDependency or reads one at a shuffle boundary). For a pipelined task set, TaskSetManager caps attempts at 1 (effectiveMaxTaskFailures) and counts failures that are normally excluded -- most importantly executor loss not caused by the app, which strands the transient output. It deliberately does NOT count the benign reasons TaskKilled (a losing duplicate attempt after another succeeded, or a deliberate kill) and TaskCommitDenied (a speculation commit race): counting those would abort the group spuriously even though the task effectively succeeded. Recovery itself is not handled here: the job aborts and the streaming layer reruns the batch from the last committed offset with fresh stage ids. Inert for non-pipelined task sets: effectiveMaxTaskFailures and the counting condition both reduce to the original expressions when isPipelined is false. Tests (TaskSetManagerSuite): a pipelined set aborts after one genuine failure; counts executor loss; does NOT abort on TaskKilled or TaskCommitDenied; a non-pipelined set is unchanged (executor loss uncounted, retries apply). (DAGSchedulerSuite): both producer and consumer task sets are marked isPipelined; regular task sets are not. Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 2 +- .../org/apache/spark/scheduler/TaskSet.scala | 8 +- .../spark/scheduler/TaskSetManager.scala | 42 ++++++- .../spark/scheduler/DAGSchedulerSuite.scala | 73 +++++++++++ .../spark/scheduler/TaskSetManagerSuite.scala | 119 ++++++++++++++++++ 5 files changed, 237 insertions(+), 7 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index bfd24ceb2f895..e002c2432c51f 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -2540,7 +2540,7 @@ private[spark] class DAGScheduler( taskScheduler.submitTasks(new TaskSet( tasks.toArray, stage.id, stage.latestInfo.attemptNumber(), jobId, properties, - stage.resourceProfileId, shuffleId)) + stage.resourceProfileId, shuffleId, isPipelined = isPipelinedGroupMember(stage))) } else { // Because we posted SparkListenerStageSubmitted earlier, we should mark // the stage as completed here in case there are no tasks to run diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSet.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSet.scala index 3513cb1f93764..407ebc53fcffc 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskSet.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSet.scala @@ -33,7 +33,13 @@ private[spark] class TaskSet( val priority: Int, val properties: Properties, val resourceProfileId: Int, - val shuffleId: Option[Int]) { + val shuffleId: Option[Int], + // True if this stage is a member of a pipelined group (connected to another stage by a + // PipelinedShuffleDependency). Such a stage's transient shuffle output cannot be re-read in + // isolation, so any task failure must fail the whole group rather than be retried per-task; + // the TaskSetManager uses this to fail fast (maxTaskFailures = 1) and to count every failure, + // including otherwise-uncounted ones like executor loss. Defaults to false. + val isPipelined: Boolean = false) { val id: String = s"$stageId.$stageAttemptId" override def toString: String = "TaskSet " + id 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 3a5f60ae19358..734500a256c19 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala @@ -70,6 +70,15 @@ private[spark] class TaskSetManager( val ser = env.closureSerializer.newInstance() val tasks = taskSet.tasks + + // A pipelined-group member reads/writes a transient shuffle that cannot be re-read in isolation, + // so a single task failure must fail the whole group (which fails the job, triggering a rerun) + // rather than be retried per-task. For such a task set we cap attempts at 1 and count every + // failure, including reasons that are normally not counted (e.g. executor loss). This is the + // native equivalent of the prototype's per-batch "any failure counts" behavior, keyed on the + // pipelined dependency (via TaskSet.isPipelined) rather than a job property. + private val effectiveMaxTaskFailures = if (taskSet.isPipelined) 1 else maxTaskFailures + private val isShuffleMapTasks = tasks(0).isInstanceOf[ShuffleMapTask] // shuffleId is only available when isShuffleMapTasks=true private val shuffleId = taskSet.shuffleId @@ -1106,16 +1115,31 @@ private[spark] class TaskSetManager( emptyTaskInfoAccumulablesAndNotifyDagScheduler(tid, tasks(index), reason, null, accumUpdates, metricPeaks) - if (!isZombie && reason.countTowardsTaskFailures) { + // A pipelined-group member's transient shuffle cannot be recovered per-task, so a GENUINE task + // failure must fail the whole group even when the reason is normally not counted -- most + // importantly executor loss not caused by the app (ExecutorLostFailure with + // exitCausedByApp=false), which strands the transient output. But we must NOT force-count + // reasons that are benign by design: TaskKilled (a losing speculative/duplicate attempt after + // another attempt succeeded, or a deliberate kill) and TaskCommitDenied (a speculation commit + // race) are not real failures, and counting them would abort the group spuriously. So for a + // pipelined set we count everything EXCEPT those two benign reasons. Non-pipelined sets are + // unchanged. + val forceCountForPipelined = taskSet.isPipelined && (reason match { + case _: TaskKilled | _: TaskCommitDenied => false + case _ => true + }) + val countTowardsTaskFailures = reason.countTowardsTaskFailures || forceCountForPipelined + if (!isZombie && countTowardsTaskFailures) { assert (null != failureReason) taskSetExcludelistHelperOpt.foreach(_.updateExcludedForFailedTask( info.host, info.executorId, index, failureReasonString)) numFailures(index) += 1 - if (numFailures(index) >= maxTaskFailures) { + if (numFailures(index) >= effectiveMaxTaskFailures) { logError(log"Task ${MDC(TASK_INDEX, index)} in stage " + taskSet.logId + - log" failed ${MDC(MAX_ATTEMPTS, maxTaskFailures)} times; aborting job") + log" failed ${MDC(MAX_ATTEMPTS, effectiveMaxTaskFailures)} times; aborting job") abort("Task %d in stage %s failed %d times, most recent failure: %s\nDriver stacktrace:" - .format(index, taskSet.id, maxTaskFailures, failureReasonString), failureException) + .format(index, taskSet.id, effectiveMaxTaskFailures, failureReasonString), + failureException) return } } @@ -1191,7 +1215,15 @@ private[spark] class TaskSetManager( // could serve the shuffle outputs or the executor lost is caused by decommission (which // can destroy the whole host). The reason is the next stage wouldn't be able to fetch the // data from this dead executor so we would need to rerun these tasks on other executors. - val maybeShuffleMapOutputLoss = isShuffleMapTasks && + // A pipelined-group member must never single-resubmit an already-successful map task: its + // transient shuffle output is not addressable and a lone producer rerun hangs the streaming + // writer in awaitTerminationAcks (SC-235532). This "Resubmitted" re-enqueue loop bypasses + // handleFailedTask, so M1.4's group-atomic abort would never see it; exclude pipelined sets + // here. A genuine executor loss still flows through the running-task loop below (iter2 -> + // handleFailedTask(ExecutorLostFailure)), which is force-counted for a pipelined set and aborts + // the whole group. Note isZombie already skips a fully-complete producer's set; this guard also + // covers a PARTIALLY-complete producer losing an executor on decommission (the SC-235532 case). + val maybeShuffleMapOutputLoss = isShuffleMapTasks && !taskSet.isPipelined && !sched.sc.shuffleDriverComponents.supportsReliableStorage() && (reason.isInstanceOf[ExecutorDecommission] || !env.blockManager.externalShuffleServiceEnabled) if (maybeShuffleMapOutputLoss && !isZombie) { diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index f43354c808652..76fe0fbe6ea8b 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6432,6 +6432,47 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } + test("pipelined shuffle: a consumer is a group member even if a REGULAR dep precedes the " + + "pipelined one in its dependency list") { + // isPipelinedGroupMember's consumer walk does `rdd.dependencies.exists { case pipelined => true; + // case regularShuffle => false; case narrow => descend }`. `exists` must continue PAST the + // regular-shuffle `false` to find a pipelined dep later in the same list. Put the regular dep + // FIRST to exercise that ordering: the consumer must still be recognized as a group member (its + // task set marked isPipelined), and co-schedule with the pipelined producer once the regular + // parent is done. + val regularProducerRdd = new MyRDD(sc, 2, Nil) + val regularDep = new ShuffleDependency(regularProducerRdd, new HashPartitioner(2)) + val pipelinedProducerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(pipelinedProducerRdd, new HashPartitioner(2)) + // Regular dep FIRST, pipelined dep SECOND. + val consumerRdd = + new MyRDD(sc, 2, List(regularDep, pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + + // Consumer waits on the regular parent; complete it so the consumer is reconsidered. + val regularStageId = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq regularProducerRdd + }.get.stageId + completeShuffleMapStageSuccessfully(regularStageId, 0, 2) + + // The consumer is now co-scheduled with the pipelined producer, and its task set must be marked + // isPipelined (which requires isPipelinedGroupMember to have found the pipelined dep despite the + // regular dep appearing first). + val consumerTaskSet = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + assert(consumerTaskSet.isPipelined, + "a consumer reading a pipelined dep must be a group member even when a regular dep is listed " + + "before it") + val pipelinedStageId = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq pipelinedProducerRdd + }.get.stageId + completeShuffleMapStageSuccessfully(pipelinedStageId, 0, 2) + complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + test("pipelined shuffle: deep chain A->B->C is submitted fully concurrently") { // A --pipelined--> B --pipelined--> C : all three co-scheduled (each edge is non-sequencing). val rddA = new MyRDD(sc, 2, Nil) @@ -7219,6 +7260,38 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } } + test("pipelined shuffle: both producer and consumer task sets are marked isPipelined") { + // The TaskSet.isPipelined flag drives group-atomic failure in the task scheduler; verify the + // DAGScheduler sets it for both members of a pipelined group (producer and consumer). + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + assert(taskSets.size === 2) + val producerTs = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq producerRdd).get + val consumerTs = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd).get + assert(producerTs.isPipelined, "the pipelined producer's task set must be marked isPipelined") + assert(consumerTs.isPipelined, "the pipelined consumer's task set must be marked isPipelined") + + completeShuffleMapStageSuccessfully(producerTs.stageId, 0, 2) + complete(consumerTs, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("regular shuffle: task sets are NOT marked isPipelined (inertness)") { + // A regular producer/consumer must not be marked isPipelined. + val producerRdd = new MyRDD(sc, 2, Nil) + val regularDep = new ShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(regularDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + assert(taskSets.head.isPipelined === false, "a regular producer must not be marked isPipelined") + completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2) + assert(taskSets(1).isPipelined === false, "a regular consumer must not be marked isPipelined") + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } } class DAGSchedulerAbortStageOffSuite extends DAGSchedulerSuite { 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 28d774eb27f7e..d72f225a95a44 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala @@ -3046,6 +3046,125 @@ class TaskSetManagerSuite "ExceptionFailure must count towards task failures (contrast)") } + // ========================================================================================== + // Pipelined-group member: group-atomic failure via job-abort (M1.4) + // ========================================================================================== + + /** A single-task TaskSet marked as a pipelined-group member. */ + private def pipelinedTaskSet(): TaskSet = { + val tasks = Array.tabulate[Task[_]](1)(i => new FakeTask(0, i, Nil)) + new TaskSet(tasks, stageId = 0, stageAttemptId = 0, priority = 0, null, + ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID, None, isPipelined = true) + } + + test("pipelined task set aborts after a single task failure (maxTaskFailures capped at 1)") { + sc = new SparkContext("local", "test") + sched = new FakeTaskScheduler(sc, ("exec1", "host1")) + val taskSet = pipelinedTaskSet() + val clock = new ManualClock + clock.advance(1) + // Even though MAX_TASK_FAILURES (4) is passed, a pipelined task set caps attempts at 1. + val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES, clock = clock) + + val offerResult = manager.resourceOffer("exec1", "host1", ANY)._1 + assert(offerResult.isDefined) + // A single ordinary task failure must abort the whole task set (group -> job), no retry. + manager.handleFailedTask(offerResult.get.taskId, TaskState.FINISHED, TaskResultLost) + assert(sched.taskSetsFailed.contains(taskSet.id), + "a pipelined task set must abort after the first task failure") + } + + 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")) + val taskSet = pipelinedTaskSet() + val clock = new ManualClock + clock.advance(1) + val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES, clock = clock) + + val offerResult = manager.resourceOffer("exec1", "host1", ANY)._1 + assert(offerResult.isDefined) + // ExecutorLostFailure with exitCausedByApp=false is normally NOT counted toward task failures, + // so a non-pipelined task set would not abort. For a pipelined task set it counts, so the + // single failure aborts the group. + manager.handleFailedTask(offerResult.get.taskId, TaskState.FINISHED, + ExecutorLostFailure("exec1", exitCausedByApp = false, reason = None)) + assert(sched.taskSetsFailed.contains(taskSet.id), + "a pipelined task set must count executor loss and abort") + } + + test("pipelined task set does NOT abort on a benign TaskKilled (losing duplicate attempt)") { + // A pipelined set with 2 tasks. One task gets a duplicate attempt; when one wins, the loser is + // killed with TaskKilled. That benign kill must NOT count as a failure and must NOT abort the + // group (the task actually succeeded). Regression for the speculative-loser abort bug. + sc = new SparkContext("local", "test") + sched = new FakeTaskScheduler(sc, ("exec1", "host1")) + val tasks = Array.tabulate[Task[_]](2)(i => new FakeTask(0, i, Nil)) + val taskSet = new TaskSet(tasks, stageId = 0, stageAttemptId = 0, priority = 0, null, + ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID, None, isPipelined = true) + val clock = new ManualClock + clock.advance(1) + val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES, clock = clock) + + // Launch task index 0 and fail it with a benign TaskKilled (as if a duplicate attempt won). + val offer = manager.resourceOffer("exec1", "host1", ANY)._1 + assert(offer.isDefined && offer.get.index === 0) + manager.handleFailedTask(offer.get.taskId, TaskState.KILLED, + TaskKilled("another attempt succeeded")) + // The group must NOT be aborted: TaskKilled is benign even for a pipelined set. + assert(!sched.taskSetsFailed.contains(taskSet.id), + "a benign TaskKilled must not abort a pipelined group") + } + + test("pipelined task set does NOT abort on TaskCommitDenied (speculation commit race)") { + sc = new SparkContext("local", "test") + sched = new FakeTaskScheduler(sc, ("exec1", "host1")) + val tasks = Array.tabulate[Task[_]](2)(i => new FakeTask(0, i, Nil)) + val taskSet = new TaskSet(tasks, stageId = 0, stageAttemptId = 0, priority = 0, null, + ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID, None, isPipelined = true) + val clock = new ManualClock + clock.advance(1) + val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES, clock = clock) + + val offer = manager.resourceOffer("exec1", "host1", ANY)._1 + assert(offer.isDefined) + manager.handleFailedTask(offer.get.taskId, TaskState.FINISHED, + TaskCommitDenied(jobID = 0, partitionID = 0, attemptNumber = 0)) + assert(!sched.taskSetsFailed.contains(taskSet.id), + "TaskCommitDenied must not abort a pipelined group") + } + + test("non-pipelined task set is unaffected: executor loss is not counted, retries still apply") { + sc = new SparkContext("local", "test") + sched = new FakeTaskScheduler(sc, ("exec1", "host1")) + val taskSet = FakeTask.createTaskSet(1) // isPipelined defaults to false + val clock = new ManualClock + clock.advance(1) + val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES, clock = clock) + + // An uncounted executor-loss failure must NOT abort a normal task set. + val offer1 = manager.resourceOffer("exec1", "host1", ANY)._1 + assert(offer1.isDefined) + manager.handleFailedTask(offer1.get.taskId, TaskState.FINISHED, + ExecutorLostFailure("exec1", exitCausedByApp = false, reason = None)) + assert(!sched.taskSetsFailed.contains(taskSet.id), + "a normal task set must not count uncaused executor loss") + + // And a normal task set still retries a counted failure up to MAX_TASK_FAILURES before abort. + (1 to MAX_TASK_FAILURES).foreach { index => + val offer = manager.resourceOffer("exec1", "host1", ANY)._1 + assert(offer.isDefined, s"expected an offer on iteration $index") + manager.handleFailedTask(offer.get.taskId, TaskState.FINISHED, TaskResultLost) + if (index < MAX_TASK_FAILURES) { + assert(!sched.taskSetsFailed.contains(taskSet.id), + s"must not abort before $MAX_TASK_FAILURES failures (iteration $index)") + } else { + assert(sched.taskSetsFailed.contains(taskSet.id), + "must abort after MAX_TASK_FAILURES counted failures") + } + } + } + } class FakeLongTasks(stageId: Int, partitionId: Int) extends FakeTask(stageId, partitionId) { From 23a1e441c48cefd122d6db9aa270ef36177133a9 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sat, 11 Jul 2026 19:20:49 +0000 Subject: [PATCH 02/39] [SPARK-XXXXX][CORE] Prevent cross-job/cross-time reuse of a pipelined producer at both layers A pipelined shuffle is transient: a once-through live stream with no retained, addressable output. Unlike a regular shuffle there is nothing for a later or concurrent job to re-read, so a pipelined producer stage must never be reused. Spark can reuse a shuffle-map stage at two independent layers, and both must be closed for a pipelined shuffle (spec S4): 1. shuffleIdToMapStage stage-object reuse: getOrCreateShuffleMapStage fails fast with PIPELINED_SHUFFLE_CROSS_JOB_REUSE if a cached pipelined producer stage would be bound to a second job. A sequential re-run is unaffected: the earlier job's cleanup drops the stage from shuffleIdToMapStage, so a later job mints a fresh producer bound only to itself (this is how an RTM micro-batch reruns). 2. MapOutputTracker-availability reuse: a pipelined shuffle is never registered with the MapOutputTracker as durable output. ShuffleMapStage tracks its completed partitions locally in a monotonic set; numAvailableOutputs and findMissingPartitions read that set for a pipelined stage. Because executor/ host loss only strips MapOutputTracker entries, it can no longer flip a completed pipelined producer to unavailable and trigger a resubmit. Together these fix SC-235532 (an RTM streaming-shuffle query hangs after executor loss: the mapper is resubmitted and its streaming writer blocks forever in awaitTerminationAcks). Two further resubmit routes for a pipelined producer are closed here: - TaskSetManager.executorLost re-enqueues an already-successful map task via the Resubmitted channel on executor decommission (when getMapOutputLocation returns None, which is always true for an unregistered pipelined shuffle). This bypasses handleFailedTask, so the group-atomic abort would never see it. Guarded with !taskSet.isPipelined; a genuine loss still routes through handleFailedTask and fails the group. - handleTaskCompletion records a pipelined success unconditionally, before the "possibly bogus epoch" check. A straggler success whose executor is already marked lost otherwise removed its partition from pendingPartitions without recording it, leaving the stage "done but not available" and driving a resubmit. (Push-based shuffle merge is disabled for a PipelinedShuffleDependency in the routing change that introduces the incremental manager, since that is where a pipelined dependency would otherwise reach the merge path; see that commit.) Inert for jobs without a pipelined dependency: every branch is gated on the dependency subtype or ShuffleMapStage.isPipelined, and the full existing DAGSchedulerSuite / TaskSetManagerSuite pass unchanged. Co-authored-by: Isaac --- .../resources/error/error-conditions.json | 6 + .../apache/spark/scheduler/DAGScheduler.scala | 41 ++- .../spark/scheduler/ShuffleMapStage.scala | 44 ++- .../spark/scheduler/DAGSchedulerSuite.scala | 289 ++++++++++++++++++ .../spark/scheduler/TaskSetManagerSuite.scala | 77 +++++ 5 files changed, 451 insertions(+), 6 deletions(-) diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index 7b31c6e253576..1bb395d1ed3ef 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -6399,6 +6399,12 @@ ], "sqlState" : "42K03" }, + "PIPELINED_SHUFFLE_CROSS_JOB_REUSE" : { + "message" : [ + "The pipelined shuffle is being reused across jobs. A pipelined (incrementally-readable) shuffle is transient and has no retained output for another job to read, so its producer stage cannot be shared across jobs." + ], + "sqlState" : "0A000" + }, "PIPELINE_DATASET_WITHOUT_FLOW" : { "message" : [ "Pipeline dataset does not have any defined flows. Please attach a query with the dataset's definition, or explicitly define at least one flow that writes to the dataset." diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index e002c2432c51f..ab76409c6cda8 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -622,6 +622,21 @@ private[spark] class DAGScheduler( firstJobId: Int): ShuffleMapStage = { shuffleIdToMapStage.get(shuffleDep.shuffleId) match { case Some(stage) => + // A pipelined shuffle is transient: it is a once-through live stream with no retained, + // addressable output, so reusing its producer stage across jobs is unsound (spec S4, + // cross-time/cross-job reuse). Reuse must be prevented explicitly -- from the scheduler's + // view a shuffle-map stage can be reused unless something forbids it. If a pipelined + // dependency's shuffleId is already bound to a stage from a different job, that is the + // forbidden cross-job reuse; fail fast (S9). (Within the same job the cached stage is the + // one we just created, so returning it is correct and not reuse.) + if (shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]] && + !stage.jobIds.contains(firstJobId)) { + throw new SparkException( + errorClass = "PIPELINED_SHUFFLE_CROSS_JOB_REUSE", + messageParameters = scala.collection.immutable.Map( + "shuffleId" -> shuffleDep.shuffleId.toString), + cause = null) + } stage case None => @@ -3081,7 +3096,31 @@ private[spark] class DAGScheduler( case smt: ShuffleMapTask => val shuffleStage = stage.asInstanceOf[ShuffleMapStage] - if (!ignoreOldTaskAttempts) { + if (shuffleStage.isPipelined) { + // A pipelined shuffle is transient and must NOT be registered with the + // MapOutputTracker as a durable, addressable output (spec S4): doing so would let + // executor/host loss strip it there, flip isAvailable to false, and resubmit the + // producer -- which then hangs its streaming writer (SC-235532). Track the completed + // partition locally and monotonically on the stage instead; the incremental shuffle + // reader discovers the producer through its own transport, not the MapOutputTracker. + // Checksum-mismatch detection does not apply (a pipelined dependency never enables + // checksum retry -- see PipelinedShuffleDependency). + // + // Record the partition and decrement pendingPartitions UNCONDITIONALLY -- outside the + // `!ignoreOldTaskAttempts` and bogus-epoch guards that gate a regular shuffle. Both + // guards exist only to avoid trusting a MapOutputTracker registration that a later + // rollback (ignoreOldTaskAttempts, from an indeterminate/rolled-back stage) or an + // executor-loss strip (bogus epoch) could invalidate. A pipelined stage never + // registers there, and its completed set is monotonic and never rolled back (a + // transient shuffle cannot be recomputed; any real group failure aborts the whole + // group, S6). Skipping the record for an "old" or "bogus" straggler would be actively + // harmful: with pendingPartitions decremented but the partition unrecorded, a dropped + // last partition leaves the stage "done but not available" -> processShuffleMapStage- + // Completion resubmits the transient producer, reopening SC-235532. An already- + // successful straggler is not a failure, so recording it is always correct. + shuffleStage.pendingPartitions -= task.partitionId + shuffleStage.addPipelinedCompletedPartition(smt.partitionId) + } else if (!ignoreOldTaskAttempts) { shuffleStage.pendingPartitions -= task.partitionId val status = event.result.asInstanceOf[MapStatus] val execId = status.location.executorId diff --git a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala index 79f7af48f102a..ebcd45ddd5cc2 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala @@ -19,7 +19,7 @@ package org.apache.spark.scheduler import scala.collection.mutable.HashSet -import org.apache.spark.{MapOutputTrackerMaster, ShuffleDependency} +import org.apache.spark.{MapOutputTrackerMaster, PipelinedShuffleDependency, ShuffleDependency} import org.apache.spark.rdd.{DeterministicLevel, RDD} import org.apache.spark.util.CallSite @@ -59,6 +59,31 @@ private[spark] class ShuffleMapStage( */ val pendingPartitions = new HashSet[Int] + /** Whether this stage produces a pipelined (incrementally-readable) shuffle. */ + val isPipelined: Boolean = shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]] + + /** + * Availability tracking for a pipelined shuffle. A pipelined shuffle is transient and is NOT + * registered with the `MapOutputTracker` as a durable, addressable output (spec S4), so its + * map-stage availability cannot be read from the tracker. Instead we track completed partitions + * here, monotonically: a partition is added when its map task succeeds and is NEVER removed on + * executor/host loss. + * + * This is the crux of avoiding the SC-235532 hang: if a pipelined shuffle's availability were + * read from the `MapOutputTracker`, losing an executor that held a completed (already-consumed) + * pipelined output would strip it there, flip `isAvailable` to false, and make the DAGScheduler + * resubmit the producer -- whose streaming writer then blocks forever waiting for termination + * acks from reducers that already finished. Keeping availability local and monotonic means + * executor loss never triggers such a resubmit; a genuine mid-group failure is handled + * group-atomically instead (S6). Unused (and empty) for a non-pipelined stage. + */ + private[this] val pipelinedCompletedPartitions = new HashSet[Int] + + /** Record a successful map task's partition as completed (pipelined stages only). */ + private[scheduler] def addPipelinedCompletedPartition(partitionId: Int): Unit = { + pipelinedCompletedPartitions += partitionId + } + override def toString: String = "ShuffleMapStage " + id /** @@ -81,7 +106,12 @@ private[spark] class ShuffleMapStage( * Number of partitions that have shuffle outputs. * When this reaches [[numPartitions]], this map stage is ready. */ - def numAvailableOutputs: Int = mapOutputTrackerMaster.getNumAvailableOutputs(shuffleDep.shuffleId) + def numAvailableOutputs: Int = { + // A pipelined shuffle is not tracked in the MapOutputTracker (see pipelinedCompletedPartitions); + // read its locally-tracked, monotonic completed set instead. + if (isPipelined) pipelinedCompletedPartitions.size + else mapOutputTrackerMaster.getNumAvailableOutputs(shuffleDep.shuffleId) + } /** * Returns true if the map stage is ready, i.e. all partitions have shuffle outputs. @@ -90,9 +120,13 @@ private[spark] class ShuffleMapStage( /** Returns the sequence of partition ids that are missing (i.e. needs to be computed). */ override def findMissingPartitions(): Seq[Int] = { - mapOutputTrackerMaster - .findMissingPartitions(shuffleDep.shuffleId) - .getOrElse(0 until numPartitions) + if (isPipelined) { + (0 until numPartitions).filterNot(pipelinedCompletedPartitions.contains) + } else { + mapOutputTrackerMaster + .findMissingPartitions(shuffleDep.shuffleId) + .getOrElse(0 until numPartitions) + } } /** diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index 76fe0fbe6ea8b..effcdcfa6c4ab 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7292,6 +7292,295 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assert(results === Map(0 -> 42, 1 -> 43)) assertDataStructuresEmpty() } + + // ========================================================================================== + // Cross-job / cross-time reuse prevention at both layers (M1.6, spec S4) -- fixes SC-235532 + // ========================================================================================== + + test("pipelined shuffle: producer availability is tracked on the stage, not the MapOutputTracker") { + // A pipelined producer's completed partitions are tracked on ShuffleMapStage (monotonic), not + // registered as durable outputs in the MapOutputTracker. Verify: after the producer's map tasks + // succeed, the stage is available WITHOUT the shuffle having map outputs in the tracker. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val shuffleId = pipelinedDep.shuffleId + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + val producerStageId = taskSets.head.stageId + val producerStage = + scheduler.stageIdToStage(producerStageId).asInstanceOf[ShuffleMapStage] + + // Complete the producer's two map tasks. + complete(taskSets.head, Seq( + (Success, makeMapStatus("hostA", 2)), + (Success, makeMapStatus("hostB", 2)))) + + // The stage is available via its local completed-partition set... + assert(producerStage.isAvailable, "pipelined producer should be available after its tasks finish") + assert(producerStage.isPipelined) + // ...but the pipelined shuffle has NO map outputs registered in the MapOutputTracker. + assert(mapOutputTracker.getNumAvailableOutputs(shuffleId) === 0, + "a pipelined shuffle must not register durable map outputs in the MapOutputTracker") + + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: numAvailableOutputs/findMissingPartitions reflect the exact " + + "incomplete partition set") { + // Pin the availability CONTRACT for a partially-complete pipelined producer, not just the + // all-done / none-done endpoints: numAvailableOutputs must be the count of completed partitions + // and findMissingPartitions must return the exact ids still missing (identity, not just size). + // Without this, a plausible refactor (revert the findMissingPartitions isPipelined branch to the + // tracker, use a bare counter, drop/take by size, or invert the filter) would silently break the + // contract yet pass every full-completion test. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + val producerTs = taskSets.head + val producerStage = + scheduler.stageIdToStage(producerTs.stageId).asInstanceOf[ShuffleMapStage] + + // Complete ONLY partition 1, leaving partition 0 missing. + runEvent(makeCompletionEvent(producerTs.tasks(1), Success, makeMapStatus("hostB", 2))) + assert(producerStage.numAvailableOutputs === 1) + assert(!producerStage.isAvailable) + assert(producerStage.findMissingPartitions() === Seq(0), + "findMissingPartitions must name the exact missing partition, not just a right-sized set") + + // Complete partition 0; now nothing is missing. + runEvent(makeCompletionEvent(producerTs.tasks(0), Success, makeMapStatus("hostA", 2))) + assert(producerStage.numAvailableOutputs === 2) + assert(producerStage.isAvailable) + assert(producerStage.findMissingPartitions() === Seq.empty) + + // Drain the consumer cleanly. + val consumerTs = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + complete(consumerTs, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a post-executor-loss straggler success must not resubmit the producer " + + "(SC-235532 bogus-epoch race)") { + // A pipelined producer's map task can succeed on an executor whose loss is already recorded + // (its StatusUpdate raced the executor-loss event). That completion hits the "possibly bogus + // epoch" branch, which for a regular shuffle simply ignores it (a healthy reattempt will + // re-register the output). But the same branch also runs `pendingPartitions -= partitionId` + // first, so if it is the last pending partition and we do NOT record it in the pipelined + // completed set, the stage looks "done but not available" and processShuffleMapStageCompletion + // resubmits the transient producer -- the exact SC-235532 hang. A pipelined stage must record + // the partition as completed even on the bogus-epoch path (its output is monotonic and the + // MapOutputTracker's executor-loss stripping does not apply to it). + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + val producerTs = taskSets.head + val producerStage = + scheduler.stageIdToStage(producerTs.stageId).asInstanceOf[ShuffleMapStage] + + // Partition 0 completes normally on executor "hostA-exec" (recorded). Note makeMapStatus(host) + // builds executorId = host + "-exec", so pass host "hostA" to get executorId "hostA-exec". + runEvent(makeCompletionEvent(producerTs.tasks(0), Success, makeMapStatus("hostA", 2))) + assert(producerStage.numAvailableOutputs === 1) + val taskSetsBefore = taskSets.size + + // Executor "hostA-exec" is lost: this records executorFailureEpoch("hostA-exec") = current epoch. + runEvent(ExecutorLost("hostA-exec", ExecutorExited(-100, false, "Container marked as failed"))) + + // Now a delayed straggler Success for partition 1 arrives FROM the lost executor "hostA-exec". + // Its task epoch is the default (-1) <= the recorded failure epoch, so it is treated as + // possibly-bogus. It must still be recorded for a pipelined stage, so the producer becomes + // available and is NOT resubmitted. + runEvent(makeCompletionEvent(producerTs.tasks(1), Success, makeMapStatus("hostA", 2))) + + assert(producerStage.isAvailable, + "a pipelined producer must be available after all partitions succeed, even via a bogus-epoch " + + "straggler") + assert(producerStage.findMissingPartitions() === Seq.empty) + assert(taskSets.size === taskSetsBefore, + "the pipelined producer must NOT be resubmitted by a post-loss straggler success (SC-235532)") + + // Consumer drains normally; no hang. + val consumerTs = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + complete(consumerTs, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a producer success with ignoreOldTaskAttempts set is still recorded " + + "(no resubmit; SC-235532)") { + // Sibling of the bogus-epoch case, for the OTHER guard that gates a regular shuffle's recording: + // ignoreOldTaskAttempts (set when a stage is rolled back, e.g. as a succeeding stage of an + // indeterminate ancestor). A pipelined producer's completed set is monotonic and never rolled + // back, so its success must be recorded even when ignoreOldTaskAttempts is true -- otherwise the + // last partition is dropped (pendingPartitions decremented, not recorded) -> "done but not + // available" -> processShuffleMapStageCompletion resubmits the transient producer (SC-235532). + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + val producerTs = taskSets.head + val producerStage = + scheduler.stageIdToStage(producerTs.stageId).asInstanceOf[ShuffleMapStage] + + runEvent(makeCompletionEvent(producerTs.tasks(0), Success, makeMapStatus("hostA", 2))) + assert(producerStage.numAvailableOutputs === 1) + val taskSetsBefore = taskSets.size + + // Force ignoreOldTaskAttempts=true for the next completion: maxAttemptIdToIgnore >= the task's + // stageAttemptId (0). For a regular stage this would drop the completion; a pipelined stage must + // still record it. + producerStage.maxAttemptIdToIgnore = Some(0) + runEvent(makeCompletionEvent(producerTs.tasks(1), Success, makeMapStatus("hostB", 2))) + + assert(producerStage.isAvailable, + "a pipelined producer success must be recorded even when ignoreOldTaskAttempts is set") + assert(producerStage.findMissingPartitions() === Seq.empty) + assert(taskSets.size === taskSetsBefore, + "the pipelined producer must NOT be resubmitted when a success arrives under " + + "ignoreOldTaskAttempts (SC-235532)") + + val consumerTs = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + complete(consumerTs, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: losing an executor does NOT flip a completed producer to unavailable " + + "or resubmit it (SC-235532)") { + // SC-235532: a completed, consumed pipelined producer whose executor is lost must NOT be + // resubmitted (which would hang the streaming writer in awaitTerminationAcks). Because pipelined + // availability is tracked on the stage (monotonic) and not the MapOutputTracker, executor loss + // cannot flip isAvailable, so the producer is not resubmitted. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + val producerStageId = taskSets.head.stageId + val producerStage = + scheduler.stageIdToStage(producerStageId).asInstanceOf[ShuffleMapStage] + // Producer completes on hostA-exec (both partitions). + complete(taskSets.head, Seq( + (Success, makeMapStatus("hostA-exec", 2)), + (Success, makeMapStatus("hostA-exec", 2)))) + assert(producerStage.isAvailable) + val taskSetsAfterProducer = taskSets.size + + // Lose the executor that ran the producer. For a regular shuffle this would strip the outputs + // and flip isAvailable -> resubmit; for a pipelined shuffle it must be inert. + runEvent(ExecutorLost("hostA-exec", ExecutorExited(-100, false, "Container marked as failed"))) + + assert(producerStage.isAvailable, + "a completed pipelined producer must remain available after executor loss (no tracker strip)") + // Pin the underlying completed set directly (not just the derived isAvailable boolean): executor + // loss must not remove any completed partition, so nothing is missing. + assert(producerStage.findMissingPartitions() === Seq.empty, + "executor loss must not strip a pipelined producer's completed partitions (monotonic)") + assert(taskSets.size === taskSetsAfterProducer, + "the pipelined producer must NOT be resubmitted on executor loss (SC-235532)") + + // The consumer completes normally; no hang, no extra producer attempt. + val consumerTaskSet = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: binding a live producer stage to a second concurrent job fails fast") { + // A pipelined shuffle is a once-through live stream with no retained output, so two concurrent + // jobs cannot share one producer stage (spec S4). While job 0 is still active its producer + // stage stays cached in shuffleIdToMapStage bound to job 0; a second concurrent job that reuses + // the SAME PipelinedShuffleDependency would bind that live stage to a second jobId, which is the + // forbidden cross-job reuse. Fail fast rather than let job 1 attach to job 0's live stream. + // (A sequential re-run is NOT this case: after job 0 finishes, cleanup drops the stage from + // shuffleIdToMapStage, so a later job gets a fresh producer bound to only its own job -- exactly + // how each RTM micro-batch reruns its producer.) + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + + // Job 0: submit but leave it active (do not complete the producer or the result stage), so the + // producer stage remains cached in shuffleIdToMapStage bound to job 0. + val firstConsumer = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(firstConsumer, Array(0, 1)) + val producerStage = + scheduler.stageIdToStage(taskSets.head.stageId).asInstanceOf[ShuffleMapStage] + assert(producerStage.jobIds === Set(0)) + val shuffleStagesBeforeJob1 = scheduler.shuffleIdToMapStage.size + val stagesBeforeJob1 = scheduler.stageIdToStage.size + + // Job 1: a second concurrent job reusing the SAME pipelined dependency -> cross-job reuse of a + // live producer stage -> fail fast. + val secondConsumer = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = {} + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(secondConsumer, Array(0, 1), listener = failListener) + assert(failure.get() != null, + "binding a live pipelined producer to a second concurrent job must fail the job") + assert(failure.get().getMessage.contains("PIPELINED_SHUFFLE_CROSS_JOB_REUSE") || + failure.get().getMessage.contains("reused across jobs"), + s"expected a cross-job-reuse error, got: ${failure.get().getMessage}") + // Job 0's producer stage was never bound to job 1, and the failed job 1 left no scheduler + // residue (the throw happened during stage creation, before any job-1 state was registered). + assert(producerStage.jobIds === Set(0)) + assert(scheduler.shuffleIdToMapStage.size === shuffleStagesBeforeJob1, + "the failed cross-job submission must not add a shuffle->stage mapping") + assert(scheduler.stageIdToStage.size === stagesBeforeJob1, + "the failed cross-job submission must not leave orphaned stages") + + // Job 0 can still complete normally -- the rejected job 1 did not corrupt its shared producer. + completeShuffleMapStageSuccessfully(producerStage.id, 0, 2) + val job0Consumer = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq firstConsumer + }.get + complete(job0Consumer, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: sequential re-run of the same producer is sound (fresh stage per job)") { + // Contrast with the concurrent case above: after a job using a pipelined dependency finishes, + // cleanup removes its producer from shuffleIdToMapStage, so re-submitting a job on the SAME + // dependency creates a fresh producer stage bound to only the new job. This is sound (it is how + // an RTM micro-batch reruns) and must NOT trip the cross-job-reuse fail-fast. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + + val consumer0 = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumer0, Array(0, 1)) + complete(taskSets(0), Seq( + (Success, makeMapStatus("hostA", 2)), + (Success, makeMapStatus("hostB", 2)))) + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + results.clear() + + // Second, sequential job on the same dependency: a fresh producer stage, no fail-fast, no hang. + // taskSets is a growing buffer across both jobs, so the second job's task sets are at index 2+. + val consumer1 = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumer1, Array(0, 1)) + complete(taskSets(2), Seq( + (Success, makeMapStatus("hostA", 2)), + (Success, makeMapStatus("hostB", 2)))) + complete(taskSets(3), Seq((Success, 4), (Success, 5))) + assert(results === Map(0 -> 4, 1 -> 5)) + assertDataStructuresEmpty() + } } class DAGSchedulerAbortStageOffSuite extends DAGSchedulerSuite { 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 d72f225a95a44..1ca6e5b7a9798 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala @@ -3134,6 +3134,83 @@ class TaskSetManagerSuite "TaskCommitDenied must not abort a pipelined group") } + test("pipelined task set does NOT single-resubmit a completed map task on executor " + + "decommission (SC-235532 channel 3)") { + // The TaskSetManager.executorLost "Resubmitted" loop re-enqueues an already-successful + // ShuffleMapTask when its executor is lost and the map output looks gone -- for a decommission, + // it looks gone whenever MapOutputTracker.getMapOutputLocation returns None. A pipelined shuffle + // is NEVER registered in the MapOutputTracker (M1.6), so getMapOutputLocation is always None and + // this loop would resubmit the lone producer task -- the exact SC-235532 hang. This loop + // bypasses handleFailedTask, so M1.4's group-atomic abort would never see it. The guard + // (!taskSet.isPipelined on maybeShuffleMapOutputLoss) must keep a pipelined set out of the loop. + // + // Scenario: a 2-task pipelined producer, one task succeeds on the to-be-lost executor while the + // other is still running (so the set is NOT a zombie and the loop is entered for a non-pipelined + // set). Decommission that executor. Assert NO Resubmitted is emitted for the completed task. + sc = new SparkContext("local", "test") + sched = new FakeTaskScheduler(sc, ("execA", "host1"), ("execB", "host2")) + sched.initialize(new FakeSchedulerBackend()) + + var resubmittedTasks = 0 + val dagScheduler = new FakeDAGScheduler(sc, sched) { + override def taskEnded( + task: Task[_], + reason: TaskEndReason, + result: Any, + accumUpdates: Seq[AccumulatorV2[_, _]], + metricPeaks: Array[Long], + taskInfo: TaskInfo): Unit = { + super.taskEnded(task, reason, result, accumUpdates, metricPeaks, taskInfo) + reason match { + case Resubmitted => resubmittedTasks += 1 + case _ => + } + } + } + sched.dagScheduler.stop() + sched.setDAGScheduler(dagScheduler) + + // Two real ShuffleMapTasks (so isShuffleMapTasks = true), marked as a pipelined-group member. + // Give each task a distinct preferred location so one lands on execA and the other on execB. + val prefs = Array( + Seq[TaskLocation](TaskLocation("host1", "execA")), + Seq[TaskLocation](TaskLocation("host2", "execB"))) + val tasks = Array.tabulate[Task[_]](2) { i => + new ShuffleMapTask(0, 0, null, new Partition { override def index: Int = i }, 1, + prefs(i), JobArtifactSet.getActiveOrDefault(sc), new Properties, null) + } + val taskSet = new TaskSet(tasks, stageId = 0, stageAttemptId = 0, priority = 0, null, + ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID, Some(0), isPipelined = true) + // A pipelined shuffle is deliberately NOT registered in the MapOutputTracker, so a decommission + // lookup of its output returns None (looks lost) -- the trigger for channel 3. + val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES) + + // Launch task 0 on execA (PROCESS_LOCAL) and task 1 on execB; complete task 0 so the set is not + // yet a zombie (task 1 still running). + val t0 = manager.resourceOffer("execA", "host1", TaskLocality.PROCESS_LOCAL)._1.get + val t1 = manager.resourceOffer("execB", "host2", TaskLocality.PROCESS_LOCAL)._1.get + assert(manager.runningTasks === 2) + assert(t0.index === 0 && t1.index === 1, "each task should land on its preferred executor") + val result0 = new DirectTaskResult[String]() { + override def value(resultSer: SerializerInstance): String = "" + } + manager.handleSuccessfulTask(t0.taskId, result0) + assert(!manager.isZombie, "the set must not be a zombie (task 1 still running)") + assert(manager.successful(t0.index)) + assert(resubmittedTasks === 0) + + // Decommission execA (which ran the completed task 0) and lose it. For a NON-pipelined set this + // would re-enqueue task 0 via Resubmitted (its output looks lost); for a pipelined set the guard + // must prevent that. + manager.executorDecommission("execA") + manager.executorLost("execA", "host1", ExecutorDecommission()) + + assert(resubmittedTasks === 0, + "a pipelined producer's completed task must NOT be single-resubmitted on decommission") + assert(manager.successful(t0.index), + "the completed task must remain successful (no Resubmitted re-enqueue)") + } + test("non-pipelined task set is unaffected: executor loss is not counted, retries still apply") { sc = new SparkContext("local", "test") sched = new FakeTaskScheduler(sc, ("exec1", "host1")) From 04bd50ef917ddab91ac5632f3c54fd13061115c3 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sat, 11 Jul 2026 19:55:57 +0000 Subject: [PATCH 03/39] [SPARK-XXXXX][CORE] Complete group-atomic failure (spec S6): fail a pipelined group on a member FetchFailed (SC-233883) A FetchFailed uniquely bypasses the group-atomic maxTaskFailures=1 lever: the base TaskSetManager marks a FetchFailed task successful, zombies the set, and returns without counting the failure (FetchFailed.countTowardsTaskFailures is false), so the whole recovery happens in the DAGScheduler's FetchFailed handler, which resubmits the map stage in isolation and recomputes serially. For a pipelined group that is a deadlock (SC-233883): the transient pipelined shuffle cannot be re-read, and members are co-scheduled, so a lone-stage resubmit is never valid (spec S6, "single-stage resubmit is disabled for group members; any member failure -- including FetchFailed -- fails the whole group"). Route a FetchFailed whose failed stage or map stage is a pipelined group member to abortStage instead: aborting the failed stage tears down the job's co-scheduled member stages via cancelRunningIndependentStages and fails the job, and the caller (the streaming batch loop) reruns the batch from scratch. The teardown is complete because a v1 pipelined group is exactly one job: the producer is registered to the same job (updateJobIdStageIdMaps walks parents), so cancelRunningIndependentStages kills/finishes it whether it is still running or already finished. This coupling to "group == one job" is guaranteed by the cross-job-reuse fail-fast (PIPELINED_SHUFFLE_CROSS_JOB_REUSE); if fan-out or cross-job sharing is ever enabled, this branch must switch to tearing down pipelinedGroupOf(failedStage) explicitly rather than relying on job-atomicity. Inert for jobs without a pipelined dependency: isPipelinedGroupMember returns false for every regular stage (it is a pipelined producer only if its shuffleDep is a PipelinedShuffleDependency, and a consumer only if it reads through one), so the original FetchFailed handling runs unchanged. The full existing DAGSchedulerSuite (including all FetchFailed / shuffle-merger tests) passes. Tests: a member FetchFailed (consumer reading the pipelined edge) fails the job with no single-stage resubmit; a producer FetchFailed reading its regular input (spec S7) also fails the group. Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 105 ++++++++++----- .../spark/scheduler/DAGSchedulerSuite.scala | 126 ++++++++++++++++++ 2 files changed, 198 insertions(+), 33 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index ab76409c6cda8..b0334ee47bbea 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -3188,6 +3188,31 @@ private[spark] class DAGScheduler( log"${MDC(STAGE_ATTEMPT_ID, task.stageAttemptId)} and there is a more recent attempt for " + log"that stage (attempt " + log"${MDC(NUM_ATTEMPT, failedStage.latestInfo.attemptNumber())}) running") + } else if (isPipelinedGroupMember(failedStage) || isPipelinedGroupMember(mapStage)) { + // Failure is group-atomic for a pipelined group (spec S6). The base scheduler handles a + // FetchFailed by resubmitting just the map stage in isolation and recomputing serially, + // but a transient pipelined shuffle cannot be re-read and its members are co-scheduled, so + // a lone-stage resubmit is never valid and would deadlock the group (SC-233883). Abort the + // whole group instead: aborting the failed stage tears down its running co-scheduled + // members and fails the job, and the caller (e.g. the streaming batch loop) reruns the + // batch from scratch. This is distinct from the maxTaskFailures=1 lever (which handles + // task failures the TaskSetManager counts): a FetchFailed is NOT counted there (the base + // TaskSetManager marks the task successful and zombies the set), so the routing to group + // failure must be enforced here. + logInfo(log"Failing pipelined group containing ${MDC(FAILED_STAGE, failedStage)} " + + log"(${MDC(FAILED_STAGE_NAME, failedStage.name)}) atomically due to a fetch failure " + + log"from ${MDC(STAGE, mapStage)} (${MDC(STAGE_NAME, mapStage.name)})") + failedStage.failedAttemptIds.add(task.stageAttemptId) + // Still unregister the failed executor's outputs, exactly as the base FetchFailed path + // does -- aborting the group tears down only THIS job's stages, but the FetchFailed is + // authoritative evidence that the executor's shuffle data is gone, and other/concurrent + // jobs sharing that executor must not keep stale MapOutputTracker entries (with an + // external shuffle service, an ExecutorLost would NOT clean these, so this is the only + // proactive channel). Safe for the pipelined shuffle itself: it registers no map outputs + // in the tracker (S6), so this can only strip regular/durable outputs. + unregisterOutputsOnFetchFailedExecutor(bmAddress, task) + abortStage(failedStage, + s"A pipelined group member failed with a fetch failure: $failureMessage", None) } else { val ignoreStageFailure = ignoreDecommissionFetchFailure && isExecutorDecommissioningOrDecommissioned(taskScheduler, bmAddress) @@ -3304,39 +3329,7 @@ private[spark] class DAGScheduler( } // TODO: mark the executor as failed only if there were lots of fetch failures on it - if (bmAddress != null) { - val externalShuffleServiceEnabled = env.blockManager.externalShuffleServiceEnabled - val isHostDecommissioned = taskScheduler - .getExecutorDecommissionState(bmAddress.executorId) - .exists(_.workerHost.isDefined) - - // Shuffle output of all executors on host `bmAddress.host` may be lost if: - // - External shuffle service is enabled, so we assume that all shuffle data on node is - // bad. - // - Host is decommissioned, thus all executors on that host will die. - val shuffleOutputOfEntireHostLost = externalShuffleServiceEnabled || - isHostDecommissioned - val hostToUnregisterOutputs = if (shuffleOutputOfEntireHostLost - && unRegisterOutputOnHostOnFetchFailure) { - Some(bmAddress.host) - } else { - // Unregister shuffle data just for one executor (we don't have any - // reason to believe shuffle data has been lost for the entire host). - None - } - removeExecutorAndUnregisterOutputs( - execId = bmAddress.executorId, - fileLost = true, - hostToUnregisterOutputs = hostToUnregisterOutputs, - maybeEpoch = Some(task.epoch), - // shuffleFileLostEpoch is ignored when a host is decommissioned because some - // decommissioned executors on that host might have been removed before this fetch - // failure and might have bumped up the shuffleFileLostEpoch. We ignore that, and - // proceed with unconditional removal of shuffle outputs from all executors on that - // host, including from those that we still haven't confirmed as lost due to heartbeat - // delays. - ignoreShuffleFileLostEpoch = isHostDecommissioned) - } + unregisterOutputsOnFetchFailedExecutor(bmAddress, task) } case failure: TaskFailedReason if task.isBarrier => @@ -3876,6 +3869,52 @@ private[spark] class DAGScheduler( maybeEpoch = None) } + /** + * On a FetchFailed, unregister the shuffle outputs of the executor (or its whole host) whose + * fetch failed, treating the FetchFailed as authoritative evidence that its shuffle data is gone. + * Extracted from the base FetchFailed handler so the pipelined-group-abort branch can also run it: + * aborting the group fails only this job's stages, but a dead executor's REGULAR outputs must + * still be cleaned up for other/concurrent jobs (with an external shuffle service, an ExecutorLost + * does not clean them, so FetchFailed is the only proactive channel). No-op when `bmAddress` is + * null. Safe for a pipelined shuffle: it registers no map outputs in the tracker, so this can only + * strip regular/durable outputs. + */ + private def unregisterOutputsOnFetchFailedExecutor( + bmAddress: BlockManagerId, task: Task[_]): Unit = { + // TODO: mark the executor as failed only if there were lots of fetch failures on it + if (bmAddress != null) { + val externalShuffleServiceEnabled = env.blockManager.externalShuffleServiceEnabled + val isHostDecommissioned = taskScheduler + .getExecutorDecommissionState(bmAddress.executorId) + .exists(_.workerHost.isDefined) + + // Shuffle output of all executors on host `bmAddress.host` may be lost if: + // - External shuffle service is enabled, so we assume that all shuffle data on node is bad. + // - Host is decommissioned, thus all executors on that host will die. + val shuffleOutputOfEntireHostLost = externalShuffleServiceEnabled || isHostDecommissioned + val hostToUnregisterOutputs = if (shuffleOutputOfEntireHostLost + && unRegisterOutputOnHostOnFetchFailure) { + Some(bmAddress.host) + } else { + // Unregister shuffle data just for one executor (we don't have any + // reason to believe shuffle data has been lost for the entire host). + None + } + removeExecutorAndUnregisterOutputs( + execId = bmAddress.executorId, + fileLost = true, + hostToUnregisterOutputs = hostToUnregisterOutputs, + maybeEpoch = Some(task.epoch), + // shuffleFileLostEpoch is ignored when a host is decommissioned because some + // decommissioned executors on that host might have been removed before this fetch + // failure and might have bumped up the shuffleFileLostEpoch. We ignore that, and + // proceed with unconditional removal of shuffle outputs from all executors on that + // host, including from those that we still haven't confirmed as lost due to heartbeat + // delays. + ignoreShuffleFileLostEpoch = isHostDecommissioned) + } + } + /** * Handles removing an executor from the BlockManagerMaster as well as unregistering shuffle * outputs for the executor or optionally its host. diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index effcdcfa6c4ab..22aebe3318684 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7581,6 +7581,132 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assert(results === Map(0 -> 4, 1 -> 5)) assertDataStructuresEmpty() } + + test("pipelined shuffle: a FetchFailed on a group member fails the group, not a single-stage " + + "resubmit (SC-233883)") { + // SC-233883: a FetchFailed must fail an RTM (pipelined) query promptly rather than trigger the + // base scheduler's single-stage resubmit -> serial recompute -> deadlock. The transient + // pipelined shuffle cannot be re-read, and members are co-scheduled, so a lone-stage resubmit is + // never valid (spec S6). Any member's FetchFailed must abort the whole group (-> job abort -> + // the caller reruns the batch). Note the base TaskSetManager does NOT count a FetchFailed (it + // marks the task successful and zombies the set), so the group-atomic maxTaskFailures=1 lever + // does not apply to FetchFailed; the routing must be enforced in the DAGScheduler. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + val producerTs = taskSets.head + val producerStage = + scheduler.stageIdToStage(producerTs.stageId).asInstanceOf[ShuffleMapStage] + + // Producer completes (its outputs are tracked locally on the pipelined stage). + complete(producerTs, Seq( + (Success, makeMapStatus("hostA", 2)), + (Success, makeMapStatus("hostB", 2)))) + val consumerTs = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + val taskSetsBeforeFetchFailure = taskSets.size + + // A consumer task hits a FetchFailed reading the pipelined shuffle. + runEvent(makeCompletionEvent( + consumerTs.tasks(0), + FetchFailed(makeBlockManagerId("hostA"), pipelinedDep.shuffleId, 0L, 0, 0, "ignored"), + null)) + + // The job must be failed (group-atomic), and no single stage may be resubmitted: no new task + // set is created, and the scheduler is not left waiting to recompute the producer in isolation. + // scheduleResubmit posts ResubmitFailedStages on a timer, so drive any pending resubmit and + // confirm nothing new is launched. + scheduler.resubmitFailedStages() + assert(failure != null, "a FetchFailed on a pipelined group member must fail the job") + assert(taskSets.size === taskSetsBeforeFetchFailure, + "a pipelined group member's FetchFailed must NOT resubmit a single stage (SC-233883)") + assert(!scheduler.runningStages.exists(_.isInstanceOf[ShuffleMapStage]), + "the pipelined producer must not be left running/resubmitted after the group fails") + sc.listenerBus.waitUntilEmpty() + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a FetchFailed on a pipelined producer reading its regular input fails " + + "the group (SC-233883)") { + // Spec S7: a pipelined group's regular INPUT edge is an ordinary shuffle. Shape: + // regularRoot --regular--> producer(pipelined) --pipelined--> consumer. + // If a producer task fails to fetch its regular input, the failedStage is the producer itself (a + // group member), so the FetchFailed must still abort the whole group rather than resubmit the + // producer in isolation. This also guards the teardown path: abortStage(failedStage=producer) + // must fail the job even though a single-stage resubmit is what the base scheduler would do. + val rootRdd = new MyRDD(sc, 2, Nil) + val regularDep = new ShuffleDependency(rootRdd, new HashPartitioner(2)) + val producerRdd = new MyRDD(sc, 2, List(regularDep), tracker = mapOutputTracker) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + + // The regular root runs first (it is a sequencing parent of the producer); complete it. + val rootTs = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq rootRdd + }.get + complete(rootTs, Seq( + (Success, makeMapStatus("hostA", 2)), + (Success, makeMapStatus("hostB", 2)))) + + // Now the producer (pipelined) and its consumer are co-scheduled. A producer task fails fetching + // the REGULAR input from the root. + val producerTs = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq producerRdd + }.get + val taskSetsBeforeFetchFailure = taskSets.size + runEvent(makeCompletionEvent( + producerTs.tasks(0), + FetchFailed(makeBlockManagerId("hostA"), regularDep.shuffleId, 0L, 0, 0, "ignored"), + null)) + + scheduler.resubmitFailedStages() + assert(failure != null, + "a FetchFailed on a pipelined producer (reading its regular input) must fail the job") + assert(taskSets.size === taskSetsBeforeFetchFailure, + "the group must not resubmit a single stage on a producer's regular-input FetchFailed") + assert(!scheduler.runningStages.exists(s => s.rdd.eq(producerRdd) || s.rdd.eq(consumerRdd)), + "neither producer nor consumer may be left running after the group fails") + sc.listenerBus.waitUntilEmpty() + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a group-aborting FetchFailed still unregisters the dead executor's " + + "regular outputs") { + // When a FetchFailed on a pipelined group member routes to group-abort, it must STILL unregister + // the failed executor's (regular, durable) shuffle outputs -- the FetchFailed is authoritative + // evidence that executor's data is gone, and other/concurrent jobs sharing it must not keep + // stale MapOutputTracker entries (with an external shuffle service an ExecutorLost would not + // clean them). Shape: regularRoot --regular--> producer(pipelined) --pipelined--> consumer; the + // producer fetch-fails its regular input from hostA-exec, aborting the group. Assert + // removeOutputsOnExecutor(hostA-exec) was still invoked (the base FetchFailed path does this; + // the group-abort path must not skip it). + val rootRdd = new MyRDD(sc, 2, Nil) + val regularDep = new ShuffleDependency(rootRdd, new HashPartitioner(2)) + val producerRdd = new MyRDD(sc, 2, List(regularDep), tracker = mapOutputTracker) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + + val rootTs = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rootRdd).get + complete(rootTs, Seq( + (Success, makeMapStatus("hostA", 2)), + (Success, makeMapStatus("hostB", 2)))) + val producerTs = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq producerRdd).get + + runEvent(makeCompletionEvent( + producerTs.tasks(0), + FetchFailed(makeBlockManagerId("hostA"), regularDep.shuffleId, 0L, 0, 0, "ignored"), + null)) + scheduler.resubmitFailedStages() + assert(failure != null, "the group must fail on the producer's regular-input FetchFailed") + // The dead executor's regular outputs are unregistered even though the group was aborted. + verify(mapOutputTracker, times(1)).removeOutputsOnExecutor("hostA-exec") + sc.listenerBus.waitUntilEmpty() + assertDataStructuresEmpty() + } } class DAGSchedulerAbortStageOffSuite extends DAGSchedulerSuite { From 6c0c3be1d94b37d1f88f6ac7c7dddd93458a79f8 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 13 Jul 2026 04:27:04 +0000 Subject: [PATCH 04/39] [SPARK-XXXXX][CORE] Test: group-atomic rerun resets per-partition commit authorization Spec S5 requires that when a pipelined group reruns after a failure, a result stage whose tasks already committed can commit again -- OutputCommitCoordinator otherwise permanently denies re-commit for a committed partition (keyed by stage id; a Success clears nothing). This is already satisfied by M1's existing mechanisms, and this test locks that in: - Group teardown runs the committed result stage through markStageAsFinished -> outputCommitCoordinator.stageEnd, clearing its committer state, so no stale authorization survives the failed attempt (option (b) in S5). - The caller's rerun is a NEW job whose stages get fresh stage ids, so the coordinator has no prior committer for them regardless (option (a) in S5). The test commits a result task in a pipelined group (registering committer state), fails the group via its producer, asserts the coordinator's committer state is cleared on teardown, then reruns the batch as a new job and asserts the rerun's result stage gets a fresh stage id and completes -- re-committing its partitions. No production change: M1's group-atomic-failure + caller-reruns-the-batch design already conforms to S5. This adds the regression test that proves it. Co-authored-by: Isaac --- .../spark/scheduler/ShuffleMapStage.scala | 4 +- .../spark/scheduler/DAGSchedulerSuite.scala | 144 +++++++++++++++--- .../spark/scheduler/TaskSetManagerSuite.scala | 15 +- 3 files changed, 133 insertions(+), 30 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala index ebcd45ddd5cc2..f17cab26320ce 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala @@ -107,8 +107,8 @@ private[spark] class ShuffleMapStage( * When this reaches [[numPartitions]], this map stage is ready. */ def numAvailableOutputs: Int = { - // A pipelined shuffle is not tracked in the MapOutputTracker (see pipelinedCompletedPartitions); - // read its locally-tracked, monotonic completed set instead. + // A pipelined shuffle is not tracked in the MapOutputTracker (see + // pipelinedCompletedPartitions); read its locally-tracked, monotonic completed set instead. if (isPipelined) pipelinedCompletedPartitions.size else mapOutputTrackerMaster.getNumAvailableOutputs(shuffleDep.shuffleId) } diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index 22aebe3318684..9791f6ddfc9dd 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6434,7 +6434,8 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti test("pipelined shuffle: a consumer is a group member even if a REGULAR dep precedes the " + "pipelined one in its dependency list") { - // isPipelinedGroupMember's consumer walk does `rdd.dependencies.exists { case pipelined => true; + // isPipelinedGroupMember's consumer walk does `rdd.dependencies.exists { case pipelined => + // true; // case regularShuffle => false; case narrow => descend }`. `exists` must continue PAST the // regular-shuffle `false` to find a pipelined dep later in the same list. Put the regular dep // FIRST to exercise that ordering: the consumer must still be recognized as a group member (its @@ -6456,13 +6457,15 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti completeShuffleMapStageSuccessfully(regularStageId, 0, 2) // The consumer is now co-scheduled with the pipelined producer, and its task set must be marked - // isPipelined (which requires isPipelinedGroupMember to have found the pipelined dep despite the + // isPipelined (which requires isPipelinedGroupMember to have found the pipelined dep despite + // the // regular dep appearing first). val consumerTaskSet = taskSets.find { ts => scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd }.get assert(consumerTaskSet.isPipelined, - "a consumer reading a pipelined dep must be a group member even when a regular dep is listed " + + "a consumer reading a pipelined dep must be a group member even when a regular dep is " + + "listed " + "before it") val pipelinedStageId = taskSets.find { ts => scheduler.stageIdToStage(ts.stageId).rdd eq pipelinedProducerRdd @@ -7268,8 +7271,10 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) submit(consumerRdd, Array(0, 1)) assert(taskSets.size === 2) - val producerTs = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq producerRdd).get - val consumerTs = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd).get + val producerTs = + taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq producerRdd).get + val consumerTs = + taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd).get assert(producerTs.isPipelined, "the pipelined producer's task set must be marked isPipelined") assert(consumerTs.isPipelined, "the pipelined consumer's task set must be marked isPipelined") @@ -7297,7 +7302,8 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti // Cross-job / cross-time reuse prevention at both layers (M1.6, spec S4) -- fixes SC-235532 // ========================================================================================== - test("pipelined shuffle: producer availability is tracked on the stage, not the MapOutputTracker") { + test("pipelined shuffle: producer availability is tracked on the stage, not the" + + "MapOutputTracker") { // A pipelined producer's completed partitions are tracked on ShuffleMapStage (monotonic), not // registered as durable outputs in the MapOutputTracker. Verify: after the producer's map tasks // succeed, the stage is available WITHOUT the shuffle having map outputs in the tracker. @@ -7316,7 +7322,8 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti (Success, makeMapStatus("hostB", 2)))) // The stage is available via its local completed-partition set... - assert(producerStage.isAvailable, "pipelined producer should be available after its tasks finish") + assert(producerStage.isAvailable, + "pipelined producer should be available after its tasks finish") assert(producerStage.isPipelined) // ...but the pipelined shuffle has NO map outputs registered in the MapOutputTracker. assert(mapOutputTracker.getNumAvailableOutputs(shuffleId) === 0, @@ -7332,8 +7339,10 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti // Pin the availability CONTRACT for a partially-complete pipelined producer, not just the // all-done / none-done endpoints: numAvailableOutputs must be the count of completed partitions // and findMissingPartitions must return the exact ids still missing (identity, not just size). - // Without this, a plausible refactor (revert the findMissingPartitions isPipelined branch to the - // tracker, use a bare counter, drop/take by size, or invert the filter) would silently break the + // Without this, a plausible refactor (revert the findMissingPartitions isPipelined branch to + // the + // tracker, use a bare counter, drop/take by size, or invert the filter) would silently break + // the // contract yet pass every full-completion test. val producerRdd = new MyRDD(sc, 2, Nil) val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) @@ -7390,7 +7399,8 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assert(producerStage.numAvailableOutputs === 1) val taskSetsBefore = taskSets.size - // Executor "hostA-exec" is lost: this records executorFailureEpoch("hostA-exec") = current epoch. + // Executor "hostA-exec" is lost: this records executorFailureEpoch("hostA-exec") = current + // epoch. runEvent(ExecutorLost("hostA-exec", ExecutorExited(-100, false, "Container marked as failed"))) // Now a delayed straggler Success for partition 1 arrives FROM the lost executor "hostA-exec". @@ -7400,7 +7410,8 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti runEvent(makeCompletionEvent(producerTs.tasks(1), Success, makeMapStatus("hostA", 2))) assert(producerStage.isAvailable, - "a pipelined producer must be available after all partitions succeed, even via a bogus-epoch " + + "a pipelined producer must be available after all partitions succeed, even via a " + + "bogus-epoch " + "straggler") assert(producerStage.findMissingPartitions() === Seq.empty) assert(taskSets.size === taskSetsBefore, @@ -7417,10 +7428,12 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti test("pipelined shuffle: a producer success with ignoreOldTaskAttempts set is still recorded " + "(no resubmit; SC-235532)") { - // Sibling of the bogus-epoch case, for the OTHER guard that gates a regular shuffle's recording: + // Sibling of the bogus-epoch case, for the OTHER guard that gates a regular shuffle's + // recording: // ignoreOldTaskAttempts (set when a stage is rolled back, e.g. as a succeeding stage of an // indeterminate ancestor). A pipelined producer's completed set is monotonic and never rolled - // back, so its success must be recorded even when ignoreOldTaskAttempts is true -- otherwise the + // back, so its success must be recorded even when ignoreOldTaskAttempts is true -- otherwise + // the // last partition is dropped (pendingPartitions decremented, not recorded) -> "done but not // available" -> processShuffleMapStageCompletion resubmits the transient producer (SC-235532). val producerRdd = new MyRDD(sc, 2, Nil) @@ -7436,7 +7449,8 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti val taskSetsBefore = taskSets.size // Force ignoreOldTaskAttempts=true for the next completion: maxAttemptIdToIgnore >= the task's - // stageAttemptId (0). For a regular stage this would drop the completion; a pipelined stage must + // stageAttemptId (0). For a regular stage this would drop the completion; a pipelined stage + // must // still record it. producerStage.maxAttemptIdToIgnore = Some(0) runEvent(makeCompletionEvent(producerTs.tasks(1), Success, makeMapStatus("hostB", 2))) @@ -7459,7 +7473,8 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti test("pipelined shuffle: losing an executor does NOT flip a completed producer to unavailable " + "or resubmit it (SC-235532)") { // SC-235532: a completed, consumed pipelined producer whose executor is lost must NOT be - // resubmitted (which would hang the streaming writer in awaitTerminationAcks). Because pipelined + // resubmitted (which would hang the streaming writer in awaitTerminationAcks). Because + // pipelined // availability is tracked on the stage (monotonic) and not the MapOutputTracker, executor loss // cannot flip isAvailable, so the producer is not resubmitted. val producerRdd = new MyRDD(sc, 2, Nil) @@ -7482,7 +7497,8 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assert(producerStage.isAvailable, "a completed pipelined producer must remain available after executor loss (no tracker strip)") - // Pin the underlying completed set directly (not just the derived isAvailable boolean): executor + // Pin the underlying completed set directly (not just the derived isAvailable boolean): + // executor // loss must not remove any completed partition, so nothing is missing. assert(producerStage.findMissingPartitions() === Seq.empty, "executor loss must not strip a pipelined producer's completed partitions (monotonic)") @@ -7502,10 +7518,12 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti // A pipelined shuffle is a once-through live stream with no retained output, so two concurrent // jobs cannot share one producer stage (spec S4). While job 0 is still active its producer // stage stays cached in shuffleIdToMapStage bound to job 0; a second concurrent job that reuses - // the SAME PipelinedShuffleDependency would bind that live stage to a second jobId, which is the + // the SAME PipelinedShuffleDependency would bind that live stage to a second jobId, which is + // the // forbidden cross-job reuse. Fail fast rather than let job 1 attach to job 0's live stream. // (A sequential re-run is NOT this case: after job 0 finishes, cleanup drops the stage from - // shuffleIdToMapStage, so a later job gets a fresh producer bound to only its own job -- exactly + // shuffleIdToMapStage, so a later job gets a fresh producer bound to only its own job -- + // exactly // how each RTM micro-batch reruns its producer.) val producerRdd = new MyRDD(sc, 2, Nil) val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) @@ -7586,7 +7604,8 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti "resubmit (SC-233883)") { // SC-233883: a FetchFailed must fail an RTM (pipelined) query promptly rather than trigger the // base scheduler's single-stage resubmit -> serial recompute -> deadlock. The transient - // pipelined shuffle cannot be re-read, and members are co-scheduled, so a lone-stage resubmit is + // pipelined shuffle cannot be re-read, and members are co-scheduled, so a lone-stage resubmit + // is // never valid (spec S6). Any member's FetchFailed must abort the whole group (-> job abort -> // the caller reruns the batch). Note the base TaskSetManager does NOT count a FetchFailed (it // marks the task successful and zombies the set), so the group-atomic maxTaskFailures=1 lever @@ -7632,7 +7651,8 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti "the group (SC-233883)") { // Spec S7: a pipelined group's regular INPUT edge is an ordinary shuffle. Shape: // regularRoot --regular--> producer(pipelined) --pipelined--> consumer. - // If a producer task fails to fetch its regular input, the failedStage is the producer itself (a + // If a producer task fails to fetch its regular input, the failedStage is the producer itself + // (a // group member), so the FetchFailed must still abort the whole group rather than resubmit the // producer in isolation. This also guards the teardown path: abortStage(failedStage=producer) // must fail the job even though a single-stage resubmit is what the base scheduler would do. @@ -7651,7 +7671,8 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti (Success, makeMapStatus("hostA", 2)), (Success, makeMapStatus("hostB", 2)))) - // Now the producer (pipelined) and its consumer are co-scheduled. A producer task fails fetching + // Now the producer (pipelined) and its consumer are co-scheduled. A producer task fails + // fetching // the REGULAR input from the root. val producerTs = taskSets.find { ts => scheduler.stageIdToStage(ts.stageId).rdd eq producerRdd @@ -7675,7 +7696,8 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti test("pipelined shuffle: a group-aborting FetchFailed still unregisters the dead executor's " + "regular outputs") { - // When a FetchFailed on a pipelined group member routes to group-abort, it must STILL unregister + // When a FetchFailed on a pipelined group member routes to group-abort, it must STILL + // unregister // the failed executor's (regular, durable) shuffle outputs -- the FetchFailed is authoritative // evidence that executor's data is gone, and other/concurrent jobs sharing it must not keep // stale MapOutputTracker entries (with an external shuffle service an ExecutorLost would not @@ -7694,7 +7716,8 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti complete(rootTs, Seq( (Success, makeMapStatus("hostA", 2)), (Success, makeMapStatus("hostB", 2)))) - val producerTs = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq producerRdd).get + val producerTs = + taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq producerRdd).get runEvent(makeCompletionEvent( producerTs.tasks(0), @@ -7707,6 +7730,81 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti sc.listenerBus.waitUntilEmpty() assertDataStructuresEmpty() } + + + // ========================================================================================== + // Group-atomic rerun resets per-partition commit authorization (M1.8, spec S5) + // ========================================================================================== + + test("pipelined shuffle: a group rerun resets per-partition commit authorization") { + // Spec S5: a PG is atomic, so a failure reruns the WHOLE group -- including a result stage + // whose + // tasks already succeeded and committed. Those committed partitions are rerun and must be + // allowed to commit again. OutputCommitCoordinator permanently denies re-commit for a committed + // partition (a Success clears nothing; keyed by stage id). M1 satisfies S5 two ways, both + // asserted here: (b) the group teardown runs the committed result stage through + // markStageAsFinished -> stageEnd, clearing its committer state (so no stale authorization + // survives); and (a) the caller's rerun is a NEW job whose stages get FRESH stage ids, so the + // coordinator has no prior committer for them regardless. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = {} + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(consumerRdd, Array(0, 1), listener = failListener) + // Locate producer/consumer defensively by RDD identity (not task-set order). + val producerTaskSet = + taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq producerRdd).get + val consumerTaskSet = + taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd).get + val firstConsumerStageId = consumerTaskSet.stageId + + // Both stages register OutputCommitCoordinator state at submission (stageStart), so the + // coordinator holds per-stage commit-authorization state that a rerun must not inherit. (This + // asserts stage state exists to be reset; it does not claim a committer was authorized -- that + // needs a live canCommit RPC, which the mock backend does not drive.) + assert(!scheduler.outputCommitCoordinator.isEmpty, + "the coordinator must hold commit-authorization state for the submitted group's stages") + + // The producer now fails -> group-atomic failure -> the whole group is torn down and the job + // fails; the caller will rerun the batch as a new job. + failed(producerTaskSet, "producer blew up") + assert(failure.get() != null, "the group must fail atomically when the producer fails") + // Teardown ran the group's stages through markStageAsFinished -> stageEnd, clearing their + // coordinator state: the commit-authorization state does not survive the failed attempt (S5). + // (isEmpty here proves stageEnd reached every stage the teardown covered.) + assert(scheduler.outputCommitCoordinator.isEmpty, + "group teardown must reset per-partition commit authorization (S5); none may survive") + assertDataStructuresEmpty() + + // The caller reruns the batch as a NEW job on the same dependency. Its stages get fresh ids, so + // the coordinator has no prior committer, and the rerun's partition 0 can commit again. Only + // the + // task sets submitted from here on belong to the rerun (earlier ones' stages were cleaned up, + // so + // look them up defensively). + val taskSetsBeforeRerun = taskSets.size + val rerunConsumer = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(rerunConsumer, Array(0, 1)) + val rerunTaskSets = taskSets.drop(taskSetsBeforeRerun) + val rerunConsumerTaskSet = rerunTaskSets.find { ts => + scheduler.stageIdToStage.get(ts.stageId).exists(_.rdd eq rerunConsumer) + }.get + val rerunProducerTaskSet = rerunTaskSets.find { ts => + scheduler.stageIdToStage.get(ts.stageId).exists(_.rdd eq producerRdd) + }.get + assert(rerunConsumerTaskSet.stageId != firstConsumerStageId, + "the rerun's result stage must get a fresh stage id (fresh coordinator state)") + complete(rerunProducerTaskSet, Seq( + (Success, makeMapStatus("hostA", 2)), + (Success, makeMapStatus("hostB", 2)))) + complete(rerunConsumerTaskSet, Seq((Success, 7), (Success, 8))) + assert(results === Map(0 -> 7, 1 -> 8), "the rerun must complete, re-committing its partitions") + assertDataStructuresEmpty() + } } class DAGSchedulerAbortStageOffSuite extends DAGSchedulerSuite { 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 1ca6e5b7a9798..734144b0ae352 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala @@ -3138,14 +3138,18 @@ class TaskSetManagerSuite "decommission (SC-235532 channel 3)") { // The TaskSetManager.executorLost "Resubmitted" loop re-enqueues an already-successful // ShuffleMapTask when its executor is lost and the map output looks gone -- for a decommission, - // it looks gone whenever MapOutputTracker.getMapOutputLocation returns None. A pipelined shuffle - // is NEVER registered in the MapOutputTracker (M1.6), so getMapOutputLocation is always None and + // it looks gone whenever MapOutputTracker.getMapOutputLocation returns None. A pipelined + // shuffle + // is NEVER registered in the MapOutputTracker (M1.6), so getMapOutputLocation is always None + // and // this loop would resubmit the lone producer task -- the exact SC-235532 hang. This loop // bypasses handleFailedTask, so M1.4's group-atomic abort would never see it. The guard - // (!taskSet.isPipelined on maybeShuffleMapOutputLoss) must keep a pipelined set out of the loop. + // (!taskSet.isPipelined on maybeShuffleMapOutputLoss) must keep a pipelined set out of the + // loop. // // Scenario: a 2-task pipelined producer, one task succeeds on the to-be-lost executor while the - // other is still running (so the set is NOT a zombie and the loop is entered for a non-pipelined + // other is still running (so the set is NOT a zombie and the loop is entered for a + // non-pipelined // set). Decommission that executor. Assert NO Resubmitted is emitted for the completed task. sc = new SparkContext("local", "test") sched = new FakeTaskScheduler(sc, ("execA", "host1"), ("execB", "host2")) @@ -3200,7 +3204,8 @@ class TaskSetManagerSuite assert(resubmittedTasks === 0) // Decommission execA (which ran the completed task 0) and lose it. For a NON-pipelined set this - // would re-enqueue task 0 via Resubmitted (its output looks lost); for a pipelined set the guard + // would re-enqueue task 0 via Resubmitted (its output looks lost); for a pipelined set the + // guard // must prevent that. manager.executorDecommission("execA") manager.executorLost("execA", "host1", ExecutorDecommission()) From 2632a33ad53c417bf9206732796ea2fd7d46bf58 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sun, 19 Jul 2026 06:20:27 +0000 Subject: [PATCH 05/39] [SPARK-XXXXX][CORE] Realign group-failure tests with the all-regular-or-all-PG restriction The M1 all-regular-or-all-PG restriction (PR3) rejects any job with a regular shuffle in a pipelined job up front, including a regular-shuffle prefix feeding a pipelined producer (regularRoot --regular--> producer(pipelined) --pipelined--> consumer). Three tests exercised that now-rejected shape: - "a consumer is a group member even if a REGULAR dep precedes the pipelined one" - "a FetchFailed on a pipelined producer reading its regular input (SC-233883)" - "a group-aborting FetchFailed still unregisters the dead executor's outputs" Remove them: the shape is out of the RTM scope (scan-of-files --pipelined--> stateful, no upstream shuffle). SC-233883's core behavior -- a member FetchFailed fails the whole group rather than resubmitting a single stage -- remains covered by the pure-shape test "a FetchFailed on a group member fails the group, not a single-stage resubmit (SC-233883)". Add a rejection test documenting that a regular-shuffle prefix feeding a pipelined producer is rejected up front. Co-authored-by: Isaac --- .../spark/scheduler/DAGSchedulerSuite.scala | 145 +++--------------- 1 file changed, 21 insertions(+), 124 deletions(-) diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index 9791f6ddfc9dd..edf7cab77f7a0 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6432,47 +6432,28 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } - test("pipelined shuffle: a consumer is a group member even if a REGULAR dep precedes the " + - "pipelined one in its dependency list") { - // isPipelinedGroupMember's consumer walk does `rdd.dependencies.exists { case pipelined => - // true; - // case regularShuffle => false; case narrow => descend }`. `exists` must continue PAST the - // regular-shuffle `false` to find a pipelined dep later in the same list. Put the regular dep - // FIRST to exercise that ordering: the consumer must still be recognized as a group member (its - // task set marked isPipelined), and co-schedule with the pipelined producer once the regular - // parent is done. - val regularProducerRdd = new MyRDD(sc, 2, Nil) - val regularDep = new ShuffleDependency(regularProducerRdd, new HashPartitioner(2)) - val pipelinedProducerRdd = new MyRDD(sc, 2, Nil) - val pipelinedDep = new PipelinedShuffleDependency(pipelinedProducerRdd, new HashPartitioner(2)) - // Regular dep FIRST, pipelined dep SECOND. - val consumerRdd = - new MyRDD(sc, 2, List(regularDep, pipelinedDep), tracker = mapOutputTracker) - submit(consumerRdd, Array(0, 1)) - - // Consumer waits on the regular parent; complete it so the consumer is reconsidered. - val regularStageId = taskSets.find { ts => - scheduler.stageIdToStage(ts.stageId).rdd eq regularProducerRdd - }.get.stageId - completeShuffleMapStageSuccessfully(regularStageId, 0, 2) + test("pipelined shuffle: a regular-shuffle prefix feeding a pipelined producer is rejected") { + // Also mixed: a regular shuffle in the PREFIX that feeds a pipelined producer + // (regularRoot --regular--> producer(pipelined) --pipelined--> consumer). Spec S7 would treat + // the regular edge as an ordinary external input to the group, but v1/M1 (RTM scope: the shape + // is scan-of-files --pipelined--> stateful, with no upstream shuffle) rejects ANY regular + // shuffle in a pipelined job up front rather than supporting a mid-DAG regular prefix. + val regularRoot = new MyRDD(sc, 2, Nil) + val regularDep = new ShuffleDependency(regularRoot, new HashPartitioner(2)) + val producerRdd = new MyRDD(sc, 2, List(regularDep), tracker = mapOutputTracker) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = {} + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(consumerRdd, Array(0, 1), listener = failListener) - // The consumer is now co-scheduled with the pipelined producer, and its task set must be marked - // isPipelined (which requires isPipelinedGroupMember to have found the pipelined dep despite - // the - // regular dep appearing first). - val consumerTaskSet = taskSets.find { ts => - scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd - }.get - assert(consumerTaskSet.isPipelined, - "a consumer reading a pipelined dep must be a group member even when a regular dep is " + - "listed " + - "before it") - val pipelinedStageId = taskSets.find { ts => - scheduler.stageIdToStage(ts.stageId).rdd eq pipelinedProducerRdd - }.get.stageId - completeShuffleMapStageSuccessfully(pipelinedStageId, 0, 2) - complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) - assert(results === Map(0 -> 42, 1 -> 43)) + assert(failure.get() != null, "a pipelined job with a regular-shuffle prefix must fail") + assert(failure.get().getMessage.contains("all-regular or all-pipelined"), + s"expected a mixed-job rejection, got: ${failure.get().getMessage}") + assert(taskSets.isEmpty, "no stage should be submitted for a rejected mixed job") assertDataStructuresEmpty() } @@ -7647,90 +7628,6 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } - test("pipelined shuffle: a FetchFailed on a pipelined producer reading its regular input fails " + - "the group (SC-233883)") { - // Spec S7: a pipelined group's regular INPUT edge is an ordinary shuffle. Shape: - // regularRoot --regular--> producer(pipelined) --pipelined--> consumer. - // If a producer task fails to fetch its regular input, the failedStage is the producer itself - // (a - // group member), so the FetchFailed must still abort the whole group rather than resubmit the - // producer in isolation. This also guards the teardown path: abortStage(failedStage=producer) - // must fail the job even though a single-stage resubmit is what the base scheduler would do. - val rootRdd = new MyRDD(sc, 2, Nil) - val regularDep = new ShuffleDependency(rootRdd, new HashPartitioner(2)) - val producerRdd = new MyRDD(sc, 2, List(regularDep), tracker = mapOutputTracker) - val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) - val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) - submit(consumerRdd, Array(0, 1)) - - // The regular root runs first (it is a sequencing parent of the producer); complete it. - val rootTs = taskSets.find { ts => - scheduler.stageIdToStage(ts.stageId).rdd eq rootRdd - }.get - complete(rootTs, Seq( - (Success, makeMapStatus("hostA", 2)), - (Success, makeMapStatus("hostB", 2)))) - - // Now the producer (pipelined) and its consumer are co-scheduled. A producer task fails - // fetching - // the REGULAR input from the root. - val producerTs = taskSets.find { ts => - scheduler.stageIdToStage(ts.stageId).rdd eq producerRdd - }.get - val taskSetsBeforeFetchFailure = taskSets.size - runEvent(makeCompletionEvent( - producerTs.tasks(0), - FetchFailed(makeBlockManagerId("hostA"), regularDep.shuffleId, 0L, 0, 0, "ignored"), - null)) - - scheduler.resubmitFailedStages() - assert(failure != null, - "a FetchFailed on a pipelined producer (reading its regular input) must fail the job") - assert(taskSets.size === taskSetsBeforeFetchFailure, - "the group must not resubmit a single stage on a producer's regular-input FetchFailed") - assert(!scheduler.runningStages.exists(s => s.rdd.eq(producerRdd) || s.rdd.eq(consumerRdd)), - "neither producer nor consumer may be left running after the group fails") - sc.listenerBus.waitUntilEmpty() - assertDataStructuresEmpty() - } - - test("pipelined shuffle: a group-aborting FetchFailed still unregisters the dead executor's " + - "regular outputs") { - // When a FetchFailed on a pipelined group member routes to group-abort, it must STILL - // unregister - // the failed executor's (regular, durable) shuffle outputs -- the FetchFailed is authoritative - // evidence that executor's data is gone, and other/concurrent jobs sharing it must not keep - // stale MapOutputTracker entries (with an external shuffle service an ExecutorLost would not - // clean them). Shape: regularRoot --regular--> producer(pipelined) --pipelined--> consumer; the - // producer fetch-fails its regular input from hostA-exec, aborting the group. Assert - // removeOutputsOnExecutor(hostA-exec) was still invoked (the base FetchFailed path does this; - // the group-abort path must not skip it). - val rootRdd = new MyRDD(sc, 2, Nil) - val regularDep = new ShuffleDependency(rootRdd, new HashPartitioner(2)) - val producerRdd = new MyRDD(sc, 2, List(regularDep), tracker = mapOutputTracker) - val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) - val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) - submit(consumerRdd, Array(0, 1)) - - val rootTs = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rootRdd).get - complete(rootTs, Seq( - (Success, makeMapStatus("hostA", 2)), - (Success, makeMapStatus("hostB", 2)))) - val producerTs = - taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq producerRdd).get - - runEvent(makeCompletionEvent( - producerTs.tasks(0), - FetchFailed(makeBlockManagerId("hostA"), regularDep.shuffleId, 0L, 0, 0, "ignored"), - null)) - scheduler.resubmitFailedStages() - assert(failure != null, "the group must fail on the producer's regular-input FetchFailed") - // The dead executor's regular outputs are unregistered even though the group was aborted. - verify(mapOutputTracker, times(1)).removeOutputsOnExecutor("hostA-exec") - sc.listenerBus.waitUntilEmpty() - assertDataStructuresEmpty() - } - // ========================================================================================== // Group-atomic rerun resets per-partition commit authorization (M1.8, spec S5) From 7d0de62f0937e9339236cae5cc5595f50d30f8b2 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sun, 19 Jul 2026 19:17:58 +0000 Subject: [PATCH 06/39] [SPARK-XXXXX][CORE] Review round 1: make the SC-235532 executor-loss test non-vacuous; fix a spec citation From an adversarial review: - The SC-235532 executor-loss regression test completed the producer with makeMapStatus("hostA-exec"), which registers the output under executor id "hostA-exec-exec" (makeBlockManagerId appends "-exec"), while the test loses executor "hostA-exec" -- so removeOutputsOnExecutor never matched and the test passed even with the ShuffleMapStage pipelined-availability fix reverted (it did not guard its named regression). Use makeMapStatus("hostA") so the lost executor id matches the registered output. Verified: the test now FAILS when the fix is reverted. - Fix a spec citation in the FetchFailed group-abort comment: "registers no map outputs in the tracker" is a spec-S4 (transient / no durable output) fact, cited as S4 everywhere else; it was mistakenly attributed to S6 (the group-atomic failure section). Co-authored-by: Isaac --- .../org/apache/spark/scheduler/DAGScheduler.scala | 2 +- .../apache/spark/scheduler/DAGSchedulerSuite.scala | 12 +++++++----- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index b0334ee47bbea..a9f1012382070 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -3209,7 +3209,7 @@ private[spark] class DAGScheduler( // jobs sharing that executor must not keep stale MapOutputTracker entries (with an // external shuffle service, an ExecutorLost would NOT clean these, so this is the only // proactive channel). Safe for the pipelined shuffle itself: it registers no map outputs - // in the tracker (S6), so this can only strip regular/durable outputs. + // in the tracker (S4), so this can only strip regular/durable outputs. unregisterOutputsOnFetchFailedExecutor(bmAddress, task) abortStage(failedStage, s"A pipelined group member failed with a fetch failure: $failureMessage", None) diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index edf7cab77f7a0..2d2363375d4d8 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7465,15 +7465,17 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti val producerStageId = taskSets.head.stageId val producerStage = scheduler.stageIdToStage(producerStageId).asInstanceOf[ShuffleMapStage] - // Producer completes on hostA-exec (both partitions). + // Producer completes on hostA (both partitions). makeMapStatus("hostA") registers the output + // under executor id "hostA-exec" (makeBlockManagerId appends "-exec"), so the ExecutorLost + // below -- which loses executor id "hostA-exec" -- actually matches the registered output. complete(taskSets.head, Seq( - (Success, makeMapStatus("hostA-exec", 2)), - (Success, makeMapStatus("hostA-exec", 2)))) + (Success, makeMapStatus("hostA", 2)), + (Success, makeMapStatus("hostA", 2)))) assert(producerStage.isAvailable) val taskSetsAfterProducer = taskSets.size - // Lose the executor that ran the producer. For a regular shuffle this would strip the outputs - // and flip isAvailable -> resubmit; for a pipelined shuffle it must be inert. + // Lose the executor that ran the producer ("hostA-exec"). For a regular shuffle this would + // strip the outputs and flip isAvailable -> resubmit; for a pipelined shuffle it must be inert. runEvent(ExecutorLost("hostA-exec", ExecutorExited(-100, false, "Container marked as failed"))) assert(producerStage.isAvailable, From ac371564fbaa5efaf71c3a0535933a866876f4f4 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 06:29:46 +0000 Subject: [PATCH 07/39] [SPARK-XXXXX][CORE] Replace internal ticket/milestone references with descriptive prose Comments and test names referenced internal Databricks Jira tickets (SC-235532, SC-233883) and internal milestone tags (M1.2 / M1.4 / M1.6 / M1.8), which are meaningless to the Apache community. Replace each with a short description of the behavior it denotes (e.g. the streaming-writer resubmit hang, the group-atomic abort), keeping the technical content and dropping the opaque identifiers. No behavior change; comment/test-name only. Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 9 +++--- .../spark/scheduler/ShuffleMapStage.scala | 3 +- .../spark/scheduler/TaskSetManager.scala | 13 ++++---- .../spark/scheduler/DAGSchedulerSuite.scala | 31 ++++++++++--------- .../spark/scheduler/TaskSetManagerSuite.scala | 10 +++--- 5 files changed, 35 insertions(+), 31 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index a9f1012382070..a351cf4254aef 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -3100,7 +3100,8 @@ private[spark] class DAGScheduler( // A pipelined shuffle is transient and must NOT be registered with the // MapOutputTracker as a durable, addressable output (spec S4): doing so would let // executor/host loss strip it there, flip isAvailable to false, and resubmit the - // producer -- which then hangs its streaming writer (SC-235532). Track the completed + // producer -- which then hangs its streaming writer (it blocks forever on + // termination acks from reducers that already finished). Track the completed // partition locally and monotonically on the stage instead; the incremental shuffle // reader discovers the producer through its own transport, not the MapOutputTracker. // Checksum-mismatch detection does not apply (a pipelined dependency never enables @@ -3116,8 +3117,8 @@ private[spark] class DAGScheduler( // group, S6). Skipping the record for an "old" or "bogus" straggler would be actively // harmful: with pendingPartitions decremented but the partition unrecorded, a dropped // last partition leaves the stage "done but not available" -> processShuffleMapStage- - // Completion resubmits the transient producer, reopening SC-235532. An already- - // successful straggler is not a failure, so recording it is always correct. + // Completion resubmits the transient producer, reopening the streaming-writer hang. + // An already-successful straggler is not a failure, so recording it is correct. shuffleStage.pendingPartitions -= task.partitionId shuffleStage.addPipelinedCompletedPartition(smt.partitionId) } else if (!ignoreOldTaskAttempts) { @@ -3192,7 +3193,7 @@ private[spark] class DAGScheduler( // Failure is group-atomic for a pipelined group (spec S6). The base scheduler handles a // FetchFailed by resubmitting just the map stage in isolation and recomputing serially, // but a transient pipelined shuffle cannot be re-read and its members are co-scheduled, so - // a lone-stage resubmit is never valid and would deadlock the group (SC-233883). Abort the + // a lone-stage resubmit is never valid and would deadlock the group. Abort the // whole group instead: aborting the failed stage tears down its running co-scheduled // members and fails the job, and the caller (e.g. the streaming batch loop) reruns the // batch from scratch. This is distinct from the maxTaskFailures=1 lever (which handles diff --git a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala index f17cab26320ce..4271d87eb4604 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala @@ -69,7 +69,8 @@ private[spark] class ShuffleMapStage( * here, monotonically: a partition is added when its map task succeeds and is NEVER removed on * executor/host loss. * - * This is the crux of avoiding the SC-235532 hang: if a pipelined shuffle's availability were + * This is the crux of avoiding the streaming-writer resubmit hang: if a pipelined shuffle's + * availability were * read from the `MapOutputTracker`, losing an executor that held a completed (already-consumed) * pipelined output would strip it there, flip `isAvailable` to false, and make the DAGScheduler * resubmit the producer -- whose streaming writer then blocks forever waiting for termination 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 734500a256c19..761a0d350529c 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala @@ -1217,12 +1217,13 @@ private[spark] class TaskSetManager( // data from this dead executor so we would need to rerun these tasks on other executors. // A pipelined-group member must never single-resubmit an already-successful map task: its // transient shuffle output is not addressable and a lone producer rerun hangs the streaming - // writer in awaitTerminationAcks (SC-235532). This "Resubmitted" re-enqueue loop bypasses - // handleFailedTask, so M1.4's group-atomic abort would never see it; exclude pipelined sets - // here. A genuine executor loss still flows through the running-task loop below (iter2 -> - // handleFailedTask(ExecutorLostFailure)), which is force-counted for a pipelined set and aborts - // the whole group. Note isZombie already skips a fully-complete producer's set; this guard also - // covers a PARTIALLY-complete producer losing an executor on decommission (the SC-235532 case). + // writer in awaitTerminationAcks. This "Resubmitted" re-enqueue loop bypasses handleFailedTask, + // so the group-atomic abort (driven from handleFailedTask via maxTaskFailures=1) would never + // see it; exclude pipelined sets here. A genuine executor loss still flows through the + // running-task loop below (iter2 -> handleFailedTask(ExecutorLostFailure)), force-counted for a + // pipelined set and aborts the whole group. Note isZombie already skips a fully-complete + // producer's set; this guard also covers a PARTIALLY-complete producer losing an executor on + // decommission. val maybeShuffleMapOutputLoss = isShuffleMapTasks && !taskSet.isPipelined && !sched.sc.shuffleDriverComponents.supportsReliableStorage() && (reason.isInstanceOf[ExecutorDecommission] || !env.blockManager.externalShuffleServiceEnabled) diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index 2d2363375d4d8..508094adc478e 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7280,7 +7280,8 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } // ========================================================================================== - // Cross-job / cross-time reuse prevention at both layers (M1.6, spec S4) -- fixes SC-235532 + // Cross-job / cross-time reuse prevention at both layers (spec S4): a consumed pipelined + // producer whose executor is lost must not be resubmitted // ========================================================================================== test("pipelined shuffle: producer availability is tracked on the stage, not the" + @@ -7356,15 +7357,15 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } test("pipelined shuffle: a post-executor-loss straggler success must not resubmit the producer " + - "(SC-235532 bogus-epoch race)") { + "(bogus-epoch race)") { // A pipelined producer's map task can succeed on an executor whose loss is already recorded // (its StatusUpdate raced the executor-loss event). That completion hits the "possibly bogus // epoch" branch, which for a regular shuffle simply ignores it (a healthy reattempt will // re-register the output). But the same branch also runs `pendingPartitions -= partitionId` // first, so if it is the last pending partition and we do NOT record it in the pipelined // completed set, the stage looks "done but not available" and processShuffleMapStageCompletion - // resubmits the transient producer -- the exact SC-235532 hang. A pipelined stage must record - // the partition as completed even on the bogus-epoch path (its output is monotonic and the + // resubmits the transient producer -- the exact streaming-writer hang. A pipelined stage must + // record the partition as completed even on the bogus-epoch path (its output is monotonic and // MapOutputTracker's executor-loss stripping does not apply to it). val producerRdd = new MyRDD(sc, 2, Nil) val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) @@ -7396,7 +7397,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti "straggler") assert(producerStage.findMissingPartitions() === Seq.empty) assert(taskSets.size === taskSetsBefore, - "the pipelined producer must NOT be resubmitted by a post-loss straggler success (SC-235532)") + "the pipelined producer must NOT be resubmitted by a post-loss straggler success") // Consumer drains normally; no hang. val consumerTs = taskSets.find { ts => @@ -7408,7 +7409,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } test("pipelined shuffle: a producer success with ignoreOldTaskAttempts set is still recorded " + - "(no resubmit; SC-235532)") { + "(no resubmit)") { // Sibling of the bogus-epoch case, for the OTHER guard that gates a regular shuffle's // recording: // ignoreOldTaskAttempts (set when a stage is rolled back, e.g. as a succeeding stage of an @@ -7416,7 +7417,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti // back, so its success must be recorded even when ignoreOldTaskAttempts is true -- otherwise // the // last partition is dropped (pendingPartitions decremented, not recorded) -> "done but not - // available" -> processShuffleMapStageCompletion resubmits the transient producer (SC-235532). + // available" -> processShuffleMapStageCompletion resubmits the transient producer. val producerRdd = new MyRDD(sc, 2, Nil) val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) @@ -7441,7 +7442,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assert(producerStage.findMissingPartitions() === Seq.empty) assert(taskSets.size === taskSetsBefore, "the pipelined producer must NOT be resubmitted when a success arrives under " + - "ignoreOldTaskAttempts (SC-235532)") + "ignoreOldTaskAttempts") val consumerTs = taskSets.find { ts => scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd @@ -7452,8 +7453,8 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } test("pipelined shuffle: losing an executor does NOT flip a completed producer to unavailable " + - "or resubmit it (SC-235532)") { - // SC-235532: a completed, consumed pipelined producer whose executor is lost must NOT be + "or resubmit it") { + // A completed, consumed pipelined producer whose executor is lost must NOT be // resubmitted (which would hang the streaming writer in awaitTerminationAcks). Because // pipelined // availability is tracked on the stage (monotonic) and not the MapOutputTracker, executor loss @@ -7486,7 +7487,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assert(producerStage.findMissingPartitions() === Seq.empty, "executor loss must not strip a pipelined producer's completed partitions (monotonic)") assert(taskSets.size === taskSetsAfterProducer, - "the pipelined producer must NOT be resubmitted on executor loss (SC-235532)") + "the pipelined producer must NOT be resubmitted on executor loss") // The consumer completes normally; no hang, no extra producer attempt. val consumerTaskSet = taskSets.find { ts => @@ -7584,8 +7585,8 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } test("pipelined shuffle: a FetchFailed on a group member fails the group, not a single-stage " + - "resubmit (SC-233883)") { - // SC-233883: a FetchFailed must fail an RTM (pipelined) query promptly rather than trigger the + "resubmit") { + // A FetchFailed must fail an RTM (pipelined) query promptly rather than trigger the // base scheduler's single-stage resubmit -> serial recompute -> deadlock. The transient // pipelined shuffle cannot be re-read, and members are co-scheduled, so a lone-stage resubmit // is @@ -7623,7 +7624,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti scheduler.resubmitFailedStages() assert(failure != null, "a FetchFailed on a pipelined group member must fail the job") assert(taskSets.size === taskSetsBeforeFetchFailure, - "a pipelined group member's FetchFailed must NOT resubmit a single stage (SC-233883)") + "a pipelined group member's FetchFailed must NOT resubmit a single stage") assert(!scheduler.runningStages.exists(_.isInstanceOf[ShuffleMapStage]), "the pipelined producer must not be left running/resubmitted after the group fails") sc.listenerBus.waitUntilEmpty() @@ -7632,7 +7633,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti // ========================================================================================== - // Group-atomic rerun resets per-partition commit authorization (M1.8, spec S5) + // Group-atomic rerun resets per-partition commit authorization (spec S5) // ========================================================================================== test("pipelined shuffle: a group rerun resets per-partition commit authorization") { 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 734144b0ae352..384cf9cdc5e03 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala @@ -3047,7 +3047,7 @@ class TaskSetManagerSuite } // ========================================================================================== - // Pipelined-group member: group-atomic failure via job-abort (M1.4) + // Pipelined-group member: group-atomic failure via job-abort // ========================================================================================== /** A single-task TaskSet marked as a pipelined-group member. */ @@ -3135,15 +3135,15 @@ class TaskSetManagerSuite } test("pipelined task set does NOT single-resubmit a completed map task on executor " + - "decommission (SC-235532 channel 3)") { + "decommission") { // The TaskSetManager.executorLost "Resubmitted" loop re-enqueues an already-successful // ShuffleMapTask when its executor is lost and the map output looks gone -- for a decommission, // it looks gone whenever MapOutputTracker.getMapOutputLocation returns None. A pipelined // shuffle - // is NEVER registered in the MapOutputTracker (M1.6), so getMapOutputLocation is always None + // is NEVER registered in the MapOutputTracker, so getMapOutputLocation is always None // and - // this loop would resubmit the lone producer task -- the exact SC-235532 hang. This loop - // bypasses handleFailedTask, so M1.4's group-atomic abort would never see it. The guard + // this loop would resubmit the lone producer task -- the exact streaming-writer hang. This loop + // bypasses handleFailedTask, so the group-atomic abort would never see it. The guard // (!taskSet.isPipelined on maybeShuffleMapOutputLoss) must keep a pipelined set out of the // loop. // From e024cef2bdea838560a662acf7406303a9647b81 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 06:54:19 +0000 Subject: [PATCH 08/39] [SPARK-XXXXX][CORE] Gate the per-submit pipelined-group check on a cheap per-job flag isPipelinedGroupMember(stage), used to tag TaskSet.isPipelined on every submitMissingTasks, walks the stage's RDD graph -- pure overhead for a regular job that has no pipelined dependency. Record whether a job uses a PipelinedShuffleDependency once at submission (ActiveJob.hasPipelinedDependency, from the hasPipelined already computed in handleJobSubmitted) and short-circuit the walk when it is false. No behavior change for a pipelined job (the flag is true, so the membership check runs exactly as before); regular jobs now skip the walk entirely. Co-authored-by: Isaac --- .../scala/org/apache/spark/scheduler/ActiveJob.scala | 8 ++++++++ .../scala/org/apache/spark/scheduler/DAGScheduler.scala | 9 ++++++++- 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/core/src/main/scala/org/apache/spark/scheduler/ActiveJob.scala b/core/src/main/scala/org/apache/spark/scheduler/ActiveJob.scala index 9876668194a84..7fa2b009fa018 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/ActiveJob.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/ActiveJob.scala @@ -63,4 +63,12 @@ private[spark] class ActiveJob( val finished = Array.fill[Boolean](numPartitions)(false) var numFinished = 0 + + /** + * Whether this job's RDD graph uses a `PipelinedShuffleDependency` (set once at job submission). + * Every pipelined-group scheduling path -- co-scheduling, deferral, and the per-submit + * `TaskSet.isPipelined` tagging -- is inert for a job without one, so this flag lets those paths + * short-circuit the group-membership graph walk for the common regular job at no cost. + */ + var hasPipelinedDependency: Boolean = false } diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index a351cf4254aef..ed63654a9d675 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -1934,6 +1934,9 @@ private[spark] class DAGScheduler( barrierJobIdToNumTasksCheckFailures.remove(jobId) val job = new ActiveJob(jobId, finalStage, callSite, listener, artifacts, properties) + // Record whether this job uses a pipelined shuffle (computed above), so the per-submit + // pipelined-group checks can short-circuit for a regular job without walking its RDD graph. + job.hasPipelinedDependency = hasPipelined clearCacheLocs() logInfo( log"Got job ${MDC(JOB_ID, job.jobId)} (${MDC(CALL_SITE_SHORT_FORM, callSite.shortForm)}) " + @@ -2553,9 +2556,13 @@ private[spark] class DAGScheduler( case _: ResultStage => None } + // Only a job that uses a pipelined shuffle can have a pipelined-group member; gate the + // group-membership graph walk on that cheap per-job flag so a regular job pays nothing here. + val isPipelined = jobIdToActiveJob.get(jobId).exists(_.hasPipelinedDependency) && + isPipelinedGroupMember(stage) taskScheduler.submitTasks(new TaskSet( tasks.toArray, stage.id, stage.latestInfo.attemptNumber(), jobId, properties, - stage.resourceProfileId, shuffleId, isPipelined = isPipelinedGroupMember(stage))) + stage.resourceProfileId, shuffleId, isPipelined = isPipelined)) } else { // Because we posted SparkListenerStageSubmitted earlier, we should mark // the stage as completed here in case there are no tasks to run From 37e3597f8c37e0635eea8aabb2e6945ba94fecf1 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 17:35:34 +0000 Subject: [PATCH 09/39] [SPARK-XXXXX][CORE] Clarify that a pipelined shuffle's MapOutputTracker entry is registered but stays empty The availability comment said a pipelined shuffle is "NOT registered with the MapOutputTracker", which could read as contradicting createShuffleMapStage (it does call registerShuffle for every ShuffleDependency, pipelined included). Clarify the precise behavior: the empty shuffle-status entry is registered to keep containsShuffle / unregisterShuffle bookkeeping uniform, but no map output is ever populated into it (registerMapOutput runs only on the non-pipelined path), so availability is read from the stage's local monotonic set instead. Comment-only. Co-authored-by: Isaac --- .../spark/scheduler/ShuffleMapStage.scala | 26 ++++++++++--------- 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala index 4271d87eb4604..c7d4590c8ede7 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala @@ -63,20 +63,22 @@ private[spark] class ShuffleMapStage( val isPipelined: Boolean = shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]] /** - * Availability tracking for a pipelined shuffle. A pipelined shuffle is transient and is NOT - * registered with the `MapOutputTracker` as a durable, addressable output (spec S4), so its - * map-stage availability cannot be read from the tracker. Instead we track completed partitions - * here, monotonically: a partition is added when its map task succeeds and is NEVER removed on - * executor/host loss. + * Availability tracking for a pipelined shuffle. A pipelined shuffle produces no durable, + * addressable map output. `createShuffleMapStage` still calls `MapOutputTracker.registerShuffle` + * for it (it is a `ShuffleDependency`, and an empty shuffle-status entry keeps the tracker's + * `containsShuffle` / `unregisterShuffle` bookkeeping uniform), but no map output is ever + * registered into that entry -- `registerMapOutput` runs only on the non-pipelined completion + * path -- so the entry stays empty. Its map-stage availability therefore cannot be read from the + * tracker; instead we track completed partitions here, monotonically: a partition is added when + * its map task succeeds and is NEVER removed on executor/host loss. * * This is the crux of avoiding the streaming-writer resubmit hang: if a pipelined shuffle's - * availability were - * read from the `MapOutputTracker`, losing an executor that held a completed (already-consumed) - * pipelined output would strip it there, flip `isAvailable` to false, and make the DAGScheduler - * resubmit the producer -- whose streaming writer then blocks forever waiting for termination - * acks from reducers that already finished. Keeping availability local and monotonic means - * executor loss never triggers such a resubmit; a genuine mid-group failure is handled - * group-atomically instead (S6). Unused (and empty) for a non-pipelined stage. + * availability were read from the `MapOutputTracker`, losing an executor that held a completed + * (already-consumed) pipelined output would strip it there, flip `isAvailable` to false, and make + * the DAGScheduler resubmit the producer -- whose streaming writer then blocks forever waiting + * for termination acks from reducers that already finished. Keeping availability local and + * monotonic means executor loss never triggers such a resubmit; a genuine mid-group failure is + * handled group-atomically instead (S6). Unused (and empty) for a non-pipelined stage. */ private[this] val pipelinedCompletedPartitions = new HashSet[Int] From e433a1ebfa1be7664c1101a1229ef295194c4c7d Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 22:18:10 +0000 Subject: [PATCH 10/39] [SPARK-XXXXX][CORE] Test the finishOnly-replay-on-torn-down-stage guard; reconcile a ShuffleMapStage comment - Regression test for the finishOnly TaskEnd guard (fixed on the scheduling PR): a deferred consumer's TaskEnd fires inline, the group is then aborted via a member FetchFailed (which removes the consumer stage), and a finishOnly replay is delivered onto the torn-down stage. Asserts the replay posts no additional TaskEnd. Lives here because it needs the group-atomic FetchFailed abort. Reverting the guard makes it fail (TaskEnd count +1). - Reconcile numAvailableOutputs's inline comment ("not tracked in the MapOutputTracker") with the field-level comment (which correctly says the entry is registered but never populated) -- both now say "outputs never populated / entry stays empty". Co-authored-by: Isaac --- .../spark/scheduler/ShuffleMapStage.scala | 4 +- .../spark/scheduler/DAGSchedulerSuite.scala | 65 +++++++++++++++++++ 2 files changed, 67 insertions(+), 2 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala index c7d4590c8ede7..0e351f5a7214a 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala @@ -110,8 +110,8 @@ private[spark] class ShuffleMapStage( * When this reaches [[numPartitions]], this map stage is ready. */ def numAvailableOutputs: Int = { - // A pipelined shuffle is not tracked in the MapOutputTracker (see - // pipelinedCompletedPartitions); read its locally-tracked, monotonic completed set instead. + // A pipelined shuffle's outputs are never populated in the MapOutputTracker (its entry stays + // empty; see pipelinedCompletedPartitions); read its locally-tracked, monotonic completed set. if (isPipelined) pipelinedCompletedPartitions.size else mapOutputTrackerMaster.getNumAvailableOutputs(shuffleDep.shuffleId) } diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index 508094adc478e..9f8c7abaceaac 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7631,6 +7631,71 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } + test("pipelined shuffle: a finishOnly replay landing on a torn-down stage does not re-post " + + "TaskEnd") { + // A deferred consumer completion fires its TaskEnd INLINE on first completion; only the + // stage/job bookkeeping is re-posted (finishOnly=true) once the producer finishes. If the + // group is torn down (a sibling aborts) between that re-post and its dequeue, the replayed + // event lands on an already-removed stage and hits the cancelled-stage guard in + // handleTaskCompletion. That guard must NOT re-post TaskEnd for a finishOnly event -- the + // per-task TaskEnd already fired, so re-posting would double-count it in listeners + // (AppStatusListener's active-task count could even go negative). This drives that exact path: + // buffer a consumer completion (TaskEnd fires once), tear the stage down, then deliver the + // finishOnly replay and assert it adds no further TaskEnd. + val taskEndCount = new java.util.concurrent.atomic.AtomicInteger(0) + val countingListener = new SparkListener { + override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = taskEndCount.incrementAndGet() + } + sc.addSparkListener(countingListener) + try { + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = results.put(index, result) + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(consumerRdd, Array(0, 1), listener = failListener) + val consumerTaskSet = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + val consumerStageId = consumerTaskSet.stageId + + // Consumer partition 0 succeeds -> its TaskEnd fires inline (count 1); bookkeeping deferred. + val replayEvent = makeCompletionEvent(consumerTaskSet.tasks(0), Success, 42) + runEvent(replayEvent) + sc.listenerBus.waitUntilEmpty() + assert(taskEndCount.get() === 1, "the consumer task's TaskEnd fires inline exactly once") + + // Tear the group down (producer FetchFailed on the consumer's other task -> group abort), + // which removes the consumer stage from stageIdToStage. (This legitimately posts its OWN + // TaskEnd for the failed task, so we snapshot the count AFTER teardown and assert the replay + // below adds nothing on top -- isolating the finishOnly replay from the abort's own events.) + runEvent(makeCompletionEvent( + consumerTaskSet.tasks(1), + FetchFailed(makeBlockManagerId("hostA"), pipelinedDep.shuffleId, 0L, 0, 0, "ignored"), + null)) + scheduler.resubmitFailedStages() + assert(failure.get() != null, "the job must fail") + assert(!scheduler.stageIdToStage.contains(consumerStageId), + "the consumer stage must have been torn down by the abort") + sc.listenerBus.waitUntilEmpty() + val countBeforeReplay = taskEndCount.get() + + // Now deliver the finishOnly replay of partition 0's completion, as the release path would + // have posted it -- but it arrives after teardown. The cancelled-stage guard must swallow it + // WITHOUT re-posting TaskEnd, so the count must not change. + runEvent(replayEvent.copy(finishOnly = true)) + sc.listenerBus.waitUntilEmpty() + assert(taskEndCount.get() === countBeforeReplay, + s"a finishOnly replay on a torn-down stage must not re-post TaskEnd; count went from " + + s"$countBeforeReplay to ${taskEndCount.get()}") + } finally { + sc.removeSparkListener(countingListener) + } + } + // ========================================================================================== // Group-atomic rerun resets per-partition commit authorization (spec S5) From 51f1ccf9e87c00b4ff642a5f27e3fbf9e4c715bd Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 22:33:55 +0000 Subject: [PATCH 11/39] [SPARK-XXXXX][CORE] Test: explicit job cancellation cleans up a buffered consumer deferral Deferral-interaction coverage. The deferral buffer (dependentStageMap) is the only mutable state this feature adds and must never outlive its job. This drives the explicit job-cancellation path (JobCancelled -> handleJobCancellation -> failJobAndIndependent- Stages), which was not directly exercised against a buffered deferral, and asserts the entry is fully cleaned up (map empty) with no buffered success applied as a result. Verified non-vacuous: neutering BOTH cleanup mechanisms -- the release-drop in cancelRunningIndependentStages (primary, fires while the producer is running) and the cleanupStateForJobAndIndependentStages sweep (backstop) -- makes the test fail with a leaked deferral entry. Co-authored-by: Isaac --- .../spark/scheduler/DAGSchedulerSuite.scala | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index 9f8c7abaceaac..34d8ff3819a15 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6484,6 +6484,35 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } + test("pipelined shuffle: an explicit job cancellation cleans up a buffered consumer deferral") { + // A buffered deferral is the only mutable state this feature adds; it must never outlive its + // job. Job cancellation goes through failJobAndIndependentStages -> + // cleanupStateForJobAndIndependentStages, which drops the entry (as a consumer key) and removes + // the stage from every other consumer's pending-producer set. Verify a consumer whose + // completion is buffered while its producer runs leaves NO deferral behind when the job is + // cancelled, and no buffered success is later applied as a result. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + val jobId = submit(consumerRdd, Array(0, 1)) + val consumerTaskSet = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + + // Consumer finishes ahead of its producer -> its completion is buffered (a deferral exists). + complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) + assert(scheduler.dependentStageMap.nonEmpty, "a deferral must exist while the producer runs") + assert(results.isEmpty, "the consumer result must be deferred, not applied yet") + + // Cancel the job. The deferral must be torn down with everything else -- no stale entry, and + // the buffered success must not be replayed as a result. + cancel(jobId) + assert(scheduler.dependentStageMap.isEmpty, + "job cancellation must clean up the buffered consumer deferral (no state outlives the job)") + assert(results.isEmpty, "a cancelled job's buffered consumer success must not be applied") + assertDataStructuresEmpty() + } + test("regular shuffle job with speculation enabled is NOT rejected (rejection path is inert)") { // The speculation fail-fast must apply only to jobs with a pipelined dependency; a plain // regular-shuffle job with speculation on runs normally. From bf81dd13994c36926c0cd23abf948a6326f9e807 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 23:09:44 +0000 Subject: [PATCH 12/39] [SPARK-XXXXX][CORE] Test: a consumer result task throwing aborts the whole pipelined group Defense-in-depth for the group-atomic failure model at the DAGScheduler layer. The maxTaskFailures=1 mechanism is covered at the TaskSetManager layer, and producer-failure teardown is covered by the deferral-drop test, but there was no DAGScheduler-level test that a CONSUMER (result) task failing tears down a still-running co-scheduled producer. This adds one: it leaves the producer running, asserts it is running, fails the consumer task set (as a maxTaskFailures=1 abort surfaces), then asserts the running producer is torn down and the job fails. Revert-checked: neutering the cancelRunningIndependentStages teardown makes it fail (producer left running). Co-authored-by: Isaac --- .../spark/scheduler/DAGSchedulerSuite.scala | 31 +++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index 34d8ff3819a15..414cf1da43c84 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7660,6 +7660,37 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } + test("pipelined shuffle: a consumer result task throwing aborts the whole group") { + // Defense-in-depth for the group-atomic model at the DAGScheduler layer: a CONSUMER (result) + // task failing must tear down the whole group, not just its own stage. Result and map tasks + // share handleFailedTask / effectiveMaxTaskFailures=1, so the first consumer-task exception + // makes the TaskSetManager abort the set (delivered here as TaskSetFailed, exactly as a + // maxTaskFailures=1 abort surfaces to the DAGScheduler), which must abort the group and tear + // down the still-running producer -- the caller then reruns the batch. Complements the + // TaskSetManager-layer maxTaskFailures=1 tests and the producer-failure drop test. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + val consumerTs = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + val producerStage = scheduler.stageIdToStage(taskSets.head.stageId) + // The producer is STILL RUNNING (co-scheduled, not yet completed) when the consumer fails -- + // this is what gives the assertion teeth: group-atomic teardown must reach a running producer. + assert(scheduler.runningStages.contains(producerStage), + "the pipelined producer must be co-scheduled and running alongside the consumer") + + // The consumer's task set aborts on the first task exception (maxTaskFailures=1 for a pipelined + // member). This must fail the whole group and tear down the still-running producer. + failed(consumerTs, "consumer result task threw") + assert(failure != null, "a consumer task failure must fail the group's job") + assert(!scheduler.runningStages.contains(producerStage), + "the still-running pipelined producer must be torn down when the consumer fails the group") + sc.listenerBus.waitUntilEmpty() + assertDataStructuresEmpty() + } + test("pipelined shuffle: a finishOnly replay landing on a torn-down stage does not re-post " + "TaskEnd") { // A deferred consumer completion fires its TaskEnd INLINE on first completion; only the From 0285b70d7a51387456a55c3279178e8e66f4c2f6 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Tue, 21 Jul 2026 07:17:10 +0000 Subject: [PATCH 13/39] [SPARK-XXXXX][CORE] Remove internal jargon from pipelined-shuffle comments Strip internal shorthand that a reader of merged Apache Spark cannot resolve -- version/milestone tags, design-doc spec section numbers, "RTM"/"PG" abbreviations, and fine-grained/coarse model terminology -- from the pipelined-shuffle failure-path comments and test comments, rewording to plain behavioral language. Comment and assertion-string text only; no behavioral change. Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 18 ++++----- .../spark/scheduler/ShuffleMapStage.scala | 2 +- .../spark/scheduler/DAGSchedulerSuite.scala | 39 +++++++++---------- 3 files changed, 28 insertions(+), 31 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index ed63654a9d675..eb0b6def479a3 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -623,11 +623,11 @@ private[spark] class DAGScheduler( shuffleIdToMapStage.get(shuffleDep.shuffleId) match { case Some(stage) => // A pipelined shuffle is transient: it is a once-through live stream with no retained, - // addressable output, so reusing its producer stage across jobs is unsound (spec S4, - // cross-time/cross-job reuse). Reuse must be prevented explicitly -- from the scheduler's - // view a shuffle-map stage can be reused unless something forbids it. If a pipelined - // dependency's shuffleId is already bound to a stage from a different job, that is the - // forbidden cross-job reuse; fail fast (S9). (Within the same job the cached stage is the + // addressable output, so reusing its producer stage across jobs is unsound (there is no + // durable output for a second job to read). Reuse must be prevented explicitly -- from the + // scheduler's view a shuffle-map stage can be reused unless something forbids it. If a + // pipelined dependency's shuffleId is already bound to a stage from a different job, that + // is the forbidden cross-job reuse; fail fast. (Within the same job the cached stage is the // one we just created, so returning it is correct and not reuse.) if (shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]] && !stage.jobIds.contains(firstJobId)) { @@ -3105,7 +3105,7 @@ private[spark] class DAGScheduler( val shuffleStage = stage.asInstanceOf[ShuffleMapStage] if (shuffleStage.isPipelined) { // A pipelined shuffle is transient and must NOT be registered with the - // MapOutputTracker as a durable, addressable output (spec S4): doing so would let + // MapOutputTracker as a durable, addressable output: doing so would let // executor/host loss strip it there, flip isAvailable to false, and resubmit the // producer -- which then hangs its streaming writer (it blocks forever on // termination acks from reducers that already finished). Track the completed @@ -3121,7 +3121,7 @@ private[spark] class DAGScheduler( // executor-loss strip (bogus epoch) could invalidate. A pipelined stage never // registers there, and its completed set is monotonic and never rolled back (a // transient shuffle cannot be recomputed; any real group failure aborts the whole - // group, S6). Skipping the record for an "old" or "bogus" straggler would be actively + // group). Skipping the record for an "old" or "bogus" straggler would be actively // harmful: with pendingPartitions decremented but the partition unrecorded, a dropped // last partition leaves the stage "done but not available" -> processShuffleMapStage- // Completion resubmits the transient producer, reopening the streaming-writer hang. @@ -3197,7 +3197,7 @@ private[spark] class DAGScheduler( log"that stage (attempt " + log"${MDC(NUM_ATTEMPT, failedStage.latestInfo.attemptNumber())}) running") } else if (isPipelinedGroupMember(failedStage) || isPipelinedGroupMember(mapStage)) { - // Failure is group-atomic for a pipelined group (spec S6). The base scheduler handles a + // Failure is group-atomic for a pipelined group. The base scheduler handles a // FetchFailed by resubmitting just the map stage in isolation and recomputing serially, // but a transient pipelined shuffle cannot be re-read and its members are co-scheduled, so // a lone-stage resubmit is never valid and would deadlock the group. Abort the @@ -3217,7 +3217,7 @@ private[spark] class DAGScheduler( // jobs sharing that executor must not keep stale MapOutputTracker entries (with an // external shuffle service, an ExecutorLost would NOT clean these, so this is the only // proactive channel). Safe for the pipelined shuffle itself: it registers no map outputs - // in the tracker (S4), so this can only strip regular/durable outputs. + // in the tracker, so this can only strip regular/durable outputs. unregisterOutputsOnFetchFailedExecutor(bmAddress, task) abortStage(failedStage, s"A pipelined group member failed with a fetch failure: $failureMessage", None) diff --git a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala index 0e351f5a7214a..94078d698cf4f 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala @@ -78,7 +78,7 @@ private[spark] class ShuffleMapStage( * the DAGScheduler resubmit the producer -- whose streaming writer then blocks forever waiting * for termination acks from reducers that already finished. Keeping availability local and * monotonic means executor loss never triggers such a resubmit; a genuine mid-group failure is - * handled group-atomically instead (S6). Unused (and empty) for a non-pipelined stage. + * handled group-atomically instead. Unused (and empty) for a non-pipelined stage. */ private[this] val pipelinedCompletedPartitions = new HashSet[Int] diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index 414cf1da43c84..e819f2d91a2e7 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6434,10 +6434,10 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti test("pipelined shuffle: a regular-shuffle prefix feeding a pipelined producer is rejected") { // Also mixed: a regular shuffle in the PREFIX that feeds a pipelined producer - // (regularRoot --regular--> producer(pipelined) --pipelined--> consumer). Spec S7 would treat - // the regular edge as an ordinary external input to the group, but v1/M1 (RTM scope: the shape - // is scan-of-files --pipelined--> stateful, with no upstream shuffle) rejects ANY regular - // shuffle in a pipelined job up front rather than supporting a mid-DAG regular prefix. + // (regularRoot --regular--> producer(pipelined) --pipelined--> consumer). Rather than treat + // the regular edge as an ordinary external input to the group and support a mid-DAG regular + // prefix, a pipelined job rejects ANY regular shuffle up front (the supported shape is + // scan-of-files --pipelined--> stateful, with no upstream shuffle). val regularRoot = new MyRDD(sc, 2, Nil) val regularDep = new ShuffleDependency(regularRoot, new HashPartitioner(2)) val producerRdd = new MyRDD(sc, 2, List(regularDep), tracker = mapOutputTracker) @@ -7309,7 +7309,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } // ========================================================================================== - // Cross-job / cross-time reuse prevention at both layers (spec S4): a consumed pipelined + // Cross-job / cross-time reuse prevention at both layers: a consumed pipelined // producer whose executor is lost must not be resubmitted // ========================================================================================== @@ -7529,15 +7529,14 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti test("pipelined shuffle: binding a live producer stage to a second concurrent job fails fast") { // A pipelined shuffle is a once-through live stream with no retained output, so two concurrent - // jobs cannot share one producer stage (spec S4). While job 0 is still active its producer + // jobs cannot share one producer stage. While job 0 is still active its producer // stage stays cached in shuffleIdToMapStage bound to job 0; a second concurrent job that reuses // the SAME PipelinedShuffleDependency would bind that live stage to a second jobId, which is // the // forbidden cross-job reuse. Fail fast rather than let job 1 attach to job 0's live stream. // (A sequential re-run is NOT this case: after job 0 finishes, cleanup drops the stage from // shuffleIdToMapStage, so a later job gets a fresh producer bound to only its own job -- - // exactly - // how each RTM micro-batch reruns its producer.) + // exactly how each streaming micro-batch reruns its producer.) val producerRdd = new MyRDD(sc, 2, Nil) val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) @@ -7587,7 +7586,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti // Contrast with the concurrent case above: after a job using a pipelined dependency finishes, // cleanup removes its producer from shuffleIdToMapStage, so re-submitting a job on the SAME // dependency creates a fresh producer stage bound to only the new job. This is sound (it is how - // an RTM micro-batch reruns) and must NOT trip the cross-job-reuse fail-fast. + // a streaming micro-batch reruns) and must NOT trip the cross-job-reuse fail-fast. val producerRdd = new MyRDD(sc, 2, Nil) val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) @@ -7615,11 +7614,10 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti test("pipelined shuffle: a FetchFailed on a group member fails the group, not a single-stage " + "resubmit") { - // A FetchFailed must fail an RTM (pipelined) query promptly rather than trigger the + // A FetchFailed must fail a pipelined (streaming) query promptly rather than trigger the // base scheduler's single-stage resubmit -> serial recompute -> deadlock. The transient // pipelined shuffle cannot be re-read, and members are co-scheduled, so a lone-stage resubmit - // is - // never valid (spec S6). Any member's FetchFailed must abort the whole group (-> job abort -> + // is never valid. Any member's FetchFailed must abort the whole group (-> job abort -> // the caller reruns the batch). Note the base TaskSetManager does NOT count a FetchFailed (it // marks the task successful and zombies the set), so the group-atomic maxTaskFailures=1 lever // does not apply to FetchFailed; the routing must be enforced in the DAGScheduler. @@ -7758,16 +7756,15 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti // ========================================================================================== - // Group-atomic rerun resets per-partition commit authorization (spec S5) + // Group-atomic rerun resets per-partition commit authorization // ========================================================================================== test("pipelined shuffle: a group rerun resets per-partition commit authorization") { - // Spec S5: a PG is atomic, so a failure reruns the WHOLE group -- including a result stage - // whose - // tasks already succeeded and committed. Those committed partitions are rerun and must be - // allowed to commit again. OutputCommitCoordinator permanently denies re-commit for a committed - // partition (a Success clears nothing; keyed by stage id). M1 satisfies S5 two ways, both - // asserted here: (b) the group teardown runs the committed result stage through + // A pipelined group is atomic, so a failure reruns the WHOLE group -- including a result + // stage whose tasks already succeeded and committed. Those committed partitions are rerun and + // must be allowed to commit again. OutputCommitCoordinator permanently denies re-commit for a + // committed partition (a Success clears nothing; keyed by stage id). The rerun stays correct + // two ways, both asserted here: (b) the group teardown runs the committed result stage through // markStageAsFinished -> stageEnd, clearing its committer state (so no stale authorization // survives); and (a) the caller's rerun is a NEW job whose stages get FRESH stage ids, so the // coordinator has no prior committer for them regardless. @@ -7799,10 +7796,10 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti failed(producerTaskSet, "producer blew up") assert(failure.get() != null, "the group must fail atomically when the producer fails") // Teardown ran the group's stages through markStageAsFinished -> stageEnd, clearing their - // coordinator state: the commit-authorization state does not survive the failed attempt (S5). + // coordinator state: the commit-authorization state does not survive the failed attempt. // (isEmpty here proves stageEnd reached every stage the teardown covered.) assert(scheduler.outputCommitCoordinator.isEmpty, - "group teardown must reset per-partition commit authorization (S5); none may survive") + "group teardown must reset per-partition commit authorization; none may survive") assertDataStructuresEmpty() // The caller reruns the batch as a NEW job on the same dependency. Its stages get fresh ids, so From af0d54149230dcb536e8e7eeda1481bbe728f88b Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Thu, 23 Jul 2026 17:15:55 +0000 Subject: [PATCH 14/39] [SPARK-XXXXX][CORE] Remove the obsolete finishOnly-replay-on-torn-down-stage test The scheduling PR now defers a pipelined consumer's whole completion event (coarse model) instead of running its per-task effects inline and replaying a finishOnly-flagged event. The finishOnly flag and its torn-down-stage guard are gone, so the regression test that exercised "a finishOnly replay landing on a torn-down stage does not re-post TaskEnd" no longer describes a reachable path: there is no inline TaskEnd and no separate finishOnly replay to land on a torn-down stage. A buffered completion is either fully replayed (producer succeeded) or dropped with its TaskEnd flushed (producer failed) -- both already covered by "deferred consumer completions are dropped when the producer fails". Remove the obsolete test. The unrelated ShuffleMapStage numAvailableOutputs comment fix from the same original commit is preserved. Co-authored-by: Isaac --- .../spark/scheduler/DAGSchedulerSuite.scala | 66 ------------------- 1 file changed, 66 deletions(-) diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index e819f2d91a2e7..7e9c74cefd131 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7689,72 +7689,6 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } - test("pipelined shuffle: a finishOnly replay landing on a torn-down stage does not re-post " + - "TaskEnd") { - // A deferred consumer completion fires its TaskEnd INLINE on first completion; only the - // stage/job bookkeeping is re-posted (finishOnly=true) once the producer finishes. If the - // group is torn down (a sibling aborts) between that re-post and its dequeue, the replayed - // event lands on an already-removed stage and hits the cancelled-stage guard in - // handleTaskCompletion. That guard must NOT re-post TaskEnd for a finishOnly event -- the - // per-task TaskEnd already fired, so re-posting would double-count it in listeners - // (AppStatusListener's active-task count could even go negative). This drives that exact path: - // buffer a consumer completion (TaskEnd fires once), tear the stage down, then deliver the - // finishOnly replay and assert it adds no further TaskEnd. - val taskEndCount = new java.util.concurrent.atomic.AtomicInteger(0) - val countingListener = new SparkListener { - override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = taskEndCount.incrementAndGet() - } - sc.addSparkListener(countingListener) - try { - val producerRdd = new MyRDD(sc, 2, Nil) - val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) - val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) - val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() - val failListener = new JobListener { - override def taskSucceeded(index: Int, result: Any): Unit = results.put(index, result) - override def jobFailed(exception: Exception): Unit = failure.set(exception) - } - submit(consumerRdd, Array(0, 1), listener = failListener) - val consumerTaskSet = taskSets.find { ts => - scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd - }.get - val consumerStageId = consumerTaskSet.stageId - - // Consumer partition 0 succeeds -> its TaskEnd fires inline (count 1); bookkeeping deferred. - val replayEvent = makeCompletionEvent(consumerTaskSet.tasks(0), Success, 42) - runEvent(replayEvent) - sc.listenerBus.waitUntilEmpty() - assert(taskEndCount.get() === 1, "the consumer task's TaskEnd fires inline exactly once") - - // Tear the group down (producer FetchFailed on the consumer's other task -> group abort), - // which removes the consumer stage from stageIdToStage. (This legitimately posts its OWN - // TaskEnd for the failed task, so we snapshot the count AFTER teardown and assert the replay - // below adds nothing on top -- isolating the finishOnly replay from the abort's own events.) - runEvent(makeCompletionEvent( - consumerTaskSet.tasks(1), - FetchFailed(makeBlockManagerId("hostA"), pipelinedDep.shuffleId, 0L, 0, 0, "ignored"), - null)) - scheduler.resubmitFailedStages() - assert(failure.get() != null, "the job must fail") - assert(!scheduler.stageIdToStage.contains(consumerStageId), - "the consumer stage must have been torn down by the abort") - sc.listenerBus.waitUntilEmpty() - val countBeforeReplay = taskEndCount.get() - - // Now deliver the finishOnly replay of partition 0's completion, as the release path would - // have posted it -- but it arrives after teardown. The cancelled-stage guard must swallow it - // WITHOUT re-posting TaskEnd, so the count must not change. - runEvent(replayEvent.copy(finishOnly = true)) - sc.listenerBus.waitUntilEmpty() - assert(taskEndCount.get() === countBeforeReplay, - s"a finishOnly replay on a torn-down stage must not re-post TaskEnd; count went from " + - s"$countBeforeReplay to ${taskEndCount.get()}") - } finally { - sc.removeSparkListener(countingListener) - } - } - - // ========================================================================================== // Group-atomic rerun resets per-partition commit authorization // ========================================================================================== From 0d435672e6cd8e8aa2388095ccd0c772f799e5da Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 13 Jul 2026 04:17:58 +0000 Subject: [PATCH 15/39] [SPARK-XXXXX][CORE] Fail fast on unsupported idioms in a pipelined group Reject, with a clear PIPELINED_SHUFFLE_UNSUPPORTED error, the pipelined-shuffle idioms a group cannot support in v1 (spec S9), so a misuse fails fast rather than silently mis-scheduling, hanging, or corrupting a group. Producer-side checks, in checkPipelinedProducerSupported (called from createShuffleMapStage when the shuffle is a PipelinedShuffleDependency), so a rejection throws during stage creation and leaves no partial scheduler state: - Barrier execution in a member stage (exposes output only after a global sync). - Dynamic resource allocation (gang admission needs a stable slot set). - Statically-indeterminate producer (its stage rollback-and-recompute recovery is moot under group-atomic failure -- a group never recomputes a single stage). - Checksum-mismatch full retry (the runtime counterpart to static indeterminism; also moot). Defensive: a PipelinedShuffleDependency does not enable it. - Push-based shuffle merge as the pipelined shuffle (finalize-then-read is the opposite of incremental reads). Defensive backstop: the dependency disables merge in its constructor. - A reliable RDD checkpoint in the member's within-stage chain (durable, lineage- truncated snapshot -> reintroduces cross-time reuse and needs a post-success recompute of the vanished transient input). Keyed on checkpointData being a ReliableRDDCheckpointData, not isCheckpointed (the write has not happened yet). Group-level check, in checkPipelinedGroupsSupportedInRDDGraph (called from handleJobSubmitted before any stage is created, alongside the speculation check, so a rejection fails the job up front with no partial state): - Fan-out: a pipelined producer feeding more than one consumer. 1:N is a supported model deferred to a later version (it needs multicast to N live readers); v1 rejects it. This also closes a group-demand undercount in the slot check when a sibling consumer stage is not yet created. Speculation and cross-job reuse (also S9) are rejected by earlier commits. The remaining S9 group-level rows -- mixed resource profiles and a regular shuffle internal to a group -- are structural invariants that do not arise for the single-profile, prefix*->PG->suffix* shapes v1 targets (a group is single-profile by construction, and groups split at regular-shuffle boundaries); they are deferred to the milestone that makes group formation first-class. Tests: barrier, indeterminate, reliable-checkpoint, and fan-out producers/groups are each rejected; regular (non-pipelined) shuffles using those same idioms are NOT rejected (inertness). Co-authored-by: Isaac --- .../resources/error/error-conditions.json | 6 + .../apache/spark/scheduler/DAGScheduler.scala | 108 +++++++++++++++++- 2 files changed, 112 insertions(+), 2 deletions(-) diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index 1bb395d1ed3ef..2e495416d67fd 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -6405,6 +6405,12 @@ ], "sqlState" : "0A000" }, + "PIPELINED_SHUFFLE_UNSUPPORTED" : { + "message" : [ + "A pipelined shuffle stage group cannot be scheduled because it uses an unsupported feature: . Pipelined (incrementally-readable) shuffles run their producer and consumer stages concurrently over a transient, once-through stream, which is incompatible with this feature in this version." + ], + "sqlState" : "0A000" + }, "PIPELINE_DATASET_WITHOUT_FLOW" : { "message" : [ "Pipeline dataset does not have any defined flows. Please attach a query with the dataset's definition, or explicitly define at least one flow that writes to the dataset." diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index eb0b6def479a3..b3dfbff75d225 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -45,7 +45,7 @@ import org.apache.spark.network.shuffle.{BlockStoreClient, MergeFinalizerListene import org.apache.spark.network.shuffle.protocol.MergeStatuses import org.apache.spark.network.util.JavaUtils import org.apache.spark.partial.{ApproximateActionListener, ApproximateEvaluator, PartialResult} -import org.apache.spark.rdd.{RDD, RDDCheckpointData} +import org.apache.spark.rdd.{DeterministicLevel, RDD, RDDCheckpointData, ReliableRDDCheckpointData} import org.apache.spark.resource.{CpuAmount, ResourceProfile, TaskResourceProfile} import org.apache.spark.resource.ResourceProfile.{CPUS, DEFAULT_RESOURCE_PROFILE_ID, EXECUTOR_CORES_LOCAL_PROPERTY, MAX_TASKS_PER_EXECUTOR_LOCAL_PROPERTY, PYSPARK_MEMORY_LOCAL_PROPERTY} import org.apache.spark.rpc.RpcTimeout @@ -686,6 +686,7 @@ private[spark] class DAGScheduler( checkBarrierStageWithDynamicAllocation(rdd) checkBarrierStageWithNumSlots(rdd, resourceProfile) checkBarrierStageWithRDDChainPattern(rdd, rdd.getNumPartitions) + checkPipelinedProducerSupported(shuffleDep) val numTasks = rdd.partitions.length val parents = getOrCreateParentStages(shuffleDeps, jobId) val id = nextStageId.getAndIncrement() @@ -709,6 +710,70 @@ private[spark] class DAGScheduler( stage } + private def pipelinedUnsupportedError(reason: String): SparkException = new SparkException( + errorClass = "PIPELINED_SHUFFLE_UNSUPPORTED", + messageParameters = scala.collection.immutable.Map("reason" -> reason), + cause = null) + + /** + * Fail-fast on producer-side idioms a pipelined shuffle cannot support in v1 (spec S9), checked + * when the producer stage is created. A pipelined shuffle runs its producer and consumer stages + * concurrently over a transient, once-through stream that a group never recomputes in isolation + * (any failure aborts the whole group, S6), so mechanisms that recompute/roll back a single stage + * are moot, and features that expose output only after a global barrier are incompatible with + * incremental reads. Rejecting here (before the stage is used) keeps a misuse from silently + * mis-scheduling. Inert for a regular ShuffleDependency. + * + * Group-level idioms (fan-out to more than one consumer, mixed resource profiles, a regular + * shuffle internal to a group) are checked separately where the group is co-scheduled, since they + * are properties of the group rather than a single producer stage. + */ + private def checkPipelinedProducerSupported(shuffleDep: ShuffleDependency[_, _, _]): Unit = { + if (!shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) { + return + } + val rdd: RDD[_] = shuffleDep.rdd + // Barrier: exposes output only after a global sync, contradicting concurrent partial reads. + if (rdd.isBarrier()) { + throw pipelinedUnsupportedError("barrier execution in a pipelined-group member stage") + } + // Dynamic resource allocation: gang admission needs a stable slot set; reclaiming executors + // from a pinned-open group can deadlock it. + if (Utils.isDynamicAllocationEnabled(sc.conf)) { + throw pipelinedUnsupportedError("dynamic resource allocation with a pipelined shuffle") + } + // Statically-indeterminate producer: its recovery is stage rollback-and-recompute, which a + // group never performs (moot under S6); reject rather than carry dead machinery. + if (rdd.outputDeterministicLevel == DeterministicLevel.INDETERMINATE) { + throw pipelinedUnsupportedError("a statically-indeterminate pipelined producer") + } + // Checksum-mismatch full retry: the runtime counterpart to static indeterminism; it rolls back + // and re-runs succeeding stages on a cross-attempt mismatch, which a group never keeps (moot). + // A PipelinedShuffleDependency does not enable it (see its definition), so this is defensive. + if (shuffleDep.checksumMismatchFullRetryEnabled) { + throw pipelinedUnsupportedError("checksum-mismatch full retry with a pipelined shuffle") + } + // Push-based shuffle merge as the pipelined shuffle: exposes output only after a post-completion + // finalize step, the opposite of incremental reads. A PipelinedShuffleDependency disables merge + // in its constructor (M1.6), so this is a defensive backstop against that being bypassed. + if (shuffleDep.shuffleMergeEnabled) { + throw pipelinedUnsupportedError("push-based shuffle merge as a pipelined shuffle") + } + // Reliable RDD checkpoint in the member's within-stage chain: writes a durable, lineage- + // truncated snapshot, which both reintroduces cross-time reuse of a transient edge (S4) and + // requires a post-success recompute of the member's now-vanished transient input. Keyed on + // checkpointData being a ReliableRDDCheckpointData (not isCheckpointed, since the write has not + // happened yet). Cache/.persist()/local checkpoint are whole-partition and ephemeral: not + // rejected. traverseParentRDDsWithinStage stops at shuffle boundaries, so only this stage's own + // chain is inspected. + val noReliableCheckpoint = traverseParentRDDsWithinStage(rdd, (r: RDD[_]) => + !r.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]])) + if (!noReliableCheckpoint) { + throw pipelinedUnsupportedError( + "a reliable RDD checkpoint in a pipelined-group member's within-stage chain") + } + } + /** * We don't support run a barrier stage with dynamic resource allocation enabled, it shall lead * to some confusing behaviors (e.g. with dynamic resource allocation enabled, it may happen that @@ -1075,6 +1140,43 @@ private[spark] class DAGScheduler( false } + /** + * Reject group-level idioms a pipelined group cannot support in v1 (spec S9), checked against the + * RDD graph BEFORE any stage is created -- so a rejection fails the job up front (via + * handleJobSubmitted's listener.jobFailed) without leaving partial scheduler state behind, exactly + * like the speculation check. Currently enforces fan-out (deferred in v1): + * + * - Fan-out: a pipelined producer feeding more than one consumer. 1:N is a supported model + * deferred to a later version (it needs multicast to N live readers); v1 rejects it. A + * PipelinedShuffleDependency's producer is `dep.rdd`; a "consumer" is any RDD that lists that + * dependency. More than one distinct consumer RDD for the same pipelined shuffle is fan-out. + * + * Inert for a job with no pipelined dependency. Throws PIPELINED_SHUFFLE_UNSUPPORTED on violation. + * + * (The other S9 group-level rows -- mixed resource profiles and a regular shuffle internal to a + * group -- are structural invariants that do not arise for the single-profile, prefix*->PG->suffix* + * shapes v1 targets: a group is single-profile by construction here, and groups are split at + * regular-shuffle boundaries (S3). Producer-side idioms -- barrier, DRA, indeterminate, checksum, + * reliable checkpoint -- are rejected in checkPipelinedProducerSupported at stage creation.) + */ + private def checkPipelinedGroupsSupportedInRDDGraph(finalRDD: RDD[_]): Unit = { + // Count distinct consumer RDDs per pipelined shuffleId across the whole RDD graph. + val consumersByShuffleId = new HashMap[Int, HashSet[Int]] + traverseRDDGraph(finalRDD) { (rdd, enqueue) => + rdd.dependencies.foreach { + case pd: PipelinedShuffleDependency[_, _, _] => + consumersByShuffleId.getOrElseUpdate(pd.shuffleId, new HashSet[Int]) += rdd.id + enqueue(pd.rdd) + case dep => + enqueue(dep.rdd) + } + } + if (consumersByShuffleId.values.exists(_.size > 1)) { + throw pipelinedUnsupportedError( + "a pipelined producer with more than one consumer (fan-out / branching)") + } + } + /** Invoke `.partitions` on the given RDD and all of its ancestors */ private def eagerlyComputePartitionsForRddAndAncestors(rdd: RDD[_]): Unit = { val startTime = System.nanoTime @@ -2095,7 +2197,9 @@ private[spark] class DAGScheduler( // rejectUnadmittablePipelinedGroup) before any member was submitted, so the group is // known to fit; just co-schedule this consumer with its running producer(s). No slot // check here -- that would re-measure capacity against a mid-flight snapshot and is - // unnecessary once admission is decided up front (gang admission). + // unnecessary once admission is decided up front (v1 gang admission). Group-level + // idiom rejection (fan-out, internal regular shuffle) already happened at job + // submission (checkPipelinedGroupsSupportedInRDDGraph + the all-PG check). logInfo(log"Submitting ${MDC(STAGE, stage)} concurrently with its running " + log"pipelined producer(s) ${MDC(MISSING_PARENT_STAGES, pipelinedMissing)}") // Record that this stage is co-scheduled with still-running pipelined producers, From d491c77464b729042dc78865209bd2abbdf0ba63 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 13 Jul 2026 04:45:36 +0000 Subject: [PATCH 16/39] [SPARK-XXXXX][CORE] Reject reliable checkpoint in a pipelined CONSUMER chain; widen group-check catch Follow-up from the M1.7 adversarial review: 1. (MEDIUM) The reliable-RDD-checkpoint rejection was scoped to a pipelined PRODUCER's within-stage chain only, but spec S9 covers any group MEMBER's chain -- including a CONSUMER, whose transient input IS the pipelined shuffle, so a reliable checkpoint there would re-read the vanished stream on recompute. Moved the reliable-checkpoint detection into checkPipelinedGroupsSupportedInRDDGraph, which runs at job submission over the whole RDD graph before any stage is created: it now walks both the producer chain (rooted at the produced RDD) and each consumer chain (rooted at a reading RDD). This also removes a partial-state leak the previous stage-creation-time consumer check would have caused (a consumer-stage throw after the producer stage had already registered). 2. (LOW) The group-check was in its own try with a SparkException-only catch, so an incidental non-Spark exception from the RDD-graph walk could escape handleJobSubmitted's job-level handler (with speculation disabled the walk is the first dependency traversal). Moved the call inside the existing createResultStage try, so any exception is handled by the same case e: Exception => jobFailed path. Not changed: fan-out keyed on shuffleId (a third, LOW review note) is correct as-is -- two distinct PipelinedShuffleDependency instances on one producer RDD are two independent 1:1 transient streams, not the multicast fan-out hazard; the check matches the spec's "more than one consumer for the same pipelined shuffle". Tests: a reliable checkpoint in a CONSUMER's chain is now rejected (new test); the producer-chain and all other rejection/inertness tests still pass. Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 113 ++++++--- .../spark/scheduler/DAGSchedulerSuite.scala | 216 +++++++++++++++++- .../spark/scheduler/TaskSetManagerSuite.scala | 15 +- 3 files changed, 308 insertions(+), 36 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index b3dfbff75d225..bd5b6e7e3dc57 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -753,25 +753,16 @@ private[spark] class DAGScheduler( if (shuffleDep.checksumMismatchFullRetryEnabled) { throw pipelinedUnsupportedError("checksum-mismatch full retry with a pipelined shuffle") } - // Push-based shuffle merge as the pipelined shuffle: exposes output only after a post-completion - // finalize step, the opposite of incremental reads. A PipelinedShuffleDependency disables merge - // in its constructor (M1.6), so this is a defensive backstop against that being bypassed. + // Push-based shuffle merge as the pipelined shuffle: exposes output only after a + // post-completion finalize step, the opposite of incremental reads. A + // PipelinedShuffleDependency disables merge in its constructor (M1.6), so this is a defensive + // backstop against that being bypassed. if (shuffleDep.shuffleMergeEnabled) { throw pipelinedUnsupportedError("push-based shuffle merge as a pipelined shuffle") } - // Reliable RDD checkpoint in the member's within-stage chain: writes a durable, lineage- - // truncated snapshot, which both reintroduces cross-time reuse of a transient edge (S4) and - // requires a post-success recompute of the member's now-vanished transient input. Keyed on - // checkpointData being a ReliableRDDCheckpointData (not isCheckpointed, since the write has not - // happened yet). Cache/.persist()/local checkpoint are whole-partition and ephemeral: not - // rejected. traverseParentRDDsWithinStage stops at shuffle boundaries, so only this stage's own - // chain is inspected. - val noReliableCheckpoint = traverseParentRDDsWithinStage(rdd, (r: RDD[_]) => - !r.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]])) - if (!noReliableCheckpoint) { - throw pipelinedUnsupportedError( - "a reliable RDD checkpoint in a pipelined-group member's within-stage chain") - } + // A reliable RDD checkpoint in a member's within-stage chain (producer OR consumer side) is + // rejected in checkPipelinedGroupsSupportedInRDDGraph, at job submission before any stage is + // created -- so a reject leaves no partial stage state and both chain sides are covered. } /** @@ -1143,29 +1134,46 @@ private[spark] class DAGScheduler( /** * Reject group-level idioms a pipelined group cannot support in v1 (spec S9), checked against the * RDD graph BEFORE any stage is created -- so a rejection fails the job up front (via - * handleJobSubmitted's listener.jobFailed) without leaving partial scheduler state behind, exactly - * like the speculation check. Currently enforces fan-out (deferred in v1): + * handleJobSubmitted's listener.jobFailed) without leaving partial scheduler state behind, + * exactly like the speculation check. * + * Inert for a job with no pipelined dependency. Throws PIPELINED_SHUFFLE_UNSUPPORTED on + * violation. Enforces: * - Fan-out: a pipelined producer feeding more than one consumer. 1:N is a supported model * deferred to a later version (it needs multicast to N live readers); v1 rejects it. A * PipelinedShuffleDependency's producer is `dep.rdd`; a "consumer" is any RDD that lists that * dependency. More than one distinct consumer RDD for the same pipelined shuffle is fan-out. - * - * Inert for a job with no pipelined dependency. Throws PIPELINED_SHUFFLE_UNSUPPORTED on violation. + * - Reliable RDD checkpoint in a group member's within-stage chain (producer OR consumer side): + * a reliable `checkpoint()` writes a durable, lineage-truncated snapshot, which both + * reintroduces cross-time reuse of a transient edge (S4) and requires a post-success recompute + * of the member's transient input -- for a consumer, that input is the vanished pipelined + * shuffle. Checked here (not at stage creation) so a reject leaves no partial stage state, and + * so BOTH the producer chain (rooted at pd.rdd) and each consumer chain (rooted at a consuming + * RDD) are covered from the whole-graph view. * * (The other S9 group-level rows -- mixed resource profiles and a regular shuffle internal to a - * group -- are structural invariants that do not arise for the single-profile, prefix*->PG->suffix* - * shapes v1 targets: a group is single-profile by construction here, and groups are split at - * regular-shuffle boundaries (S3). Producer-side idioms -- barrier, DRA, indeterminate, checksum, - * reliable checkpoint -- are rejected in checkPipelinedProducerSupported at stage creation.) + * group -- are structural invariants that do not arise for the single-profile, + * prefix*->PG->suffix* shapes v1 targets: a group is single-profile by construction here, and + * groups are split at + * regular-shuffle boundaries (S3). The remaining producer-side idioms -- barrier, DRA, + * indeterminate, checksum, push-merge -- are rejected in checkPipelinedProducerSupported at stage + * creation, where a producer-only throw leaves no partial state.) */ private def checkPipelinedGroupsSupportedInRDDGraph(finalRDD: RDD[_]): Unit = { - // Count distinct consumer RDDs per pipelined shuffleId across the whole RDD graph. + // Walk the whole RDD graph once, collecting for each pipelined shuffleId the distinct consumer + // RDDs that read it (for the fan-out check), the producer RDDs that write it (roots of producer + // member stages), and every reliably-checkpointed RDD (to locate ones inside a member stage). val consumersByShuffleId = new HashMap[Int, HashSet[Int]] + val producerRoots = new HashSet[RDD[_]] // RDDs that WRITE a pipelined shuffle + val reliablyCheckpointed = new HashSet[RDD[_]] // RDDs with a reliable checkpoint pending traverseRDDGraph(finalRDD) { (rdd, enqueue) => + if (rdd.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]])) { + reliablyCheckpointed += rdd + } rdd.dependencies.foreach { case pd: PipelinedShuffleDependency[_, _, _] => consumersByShuffleId.getOrElseUpdate(pd.shuffleId, new HashSet[Int]) += rdd.id + producerRoots += pd.rdd enqueue(pd.rdd) case dep => enqueue(dep.rdd) @@ -1175,6 +1183,47 @@ private[spark] class DAGScheduler( throw pipelinedUnsupportedError( "a pipelined producer with more than one consumer (fan-out / branching)") } + // Reject a reliable checkpoint anywhere in a pipelined-group MEMBER's within-stage chain (spec + // S9). Keyed on checkpointData being ReliableRDDCheckpointData, not isCheckpointed, since the + // write has not happened yet. Cache / .persist() / local checkpoint are whole-partition and + // ephemeral and are not rejected. A reliably-checkpointed RDD `cp` is inside a member stage + // iff: + // - PRODUCER side: cp is within some producer root's own within-stage chain (walk parents from + // the producer root, stopping at shuffle boundaries), OR + // - CONSUMER side: cp's OWN within-stage chain reads a pipelined shuffle (walk parents from + // cp, stopping at shuffle boundaries, and check whether any stopped-at boundary is + // pipelined). + // Rooting the consumer check at each checkpointed RDD (rather than at the PSD-reading RDD) is + // what makes it cover a checkpoint anywhere DOWNSTREAM in the consumer stage, not just on the + // reading RDD itself. + def chainHasReliableCheckpoint(root: RDD[_]): Boolean = + !traverseParentRDDsWithinStage(root, (r: RDD[_]) => + !r.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]])) + val offending = + // CONSUMER side: a checkpointed RDD whose own within-stage chain reads a pipelined shuffle is + // inside a consumer member stage (covers a checkpoint anywhere in that stage, not just on the + // reading RDD). PRODUCER side: a producer root's within-stage chain carries a checkpoint. + reliablyCheckpointed.exists(rddChainReadsPipelinedShuffle) || + producerRoots.exists(chainHasReliableCheckpoint) + if (offending) { + throw pipelinedUnsupportedError( + "a reliable RDD checkpoint in a pipelined-group member's within-stage chain") + } + } + + /** Whether `rdd`'s within-stage chain (parents, stopping at shuffle boundaries) reads through a + * [[PipelinedShuffleDependency]] -- i.e. `rdd` is inside a pipelined CONSUMER member stage. */ + private def rddChainReadsPipelinedShuffle(rdd: RDD[_]): Boolean = { + !traverseRDDGraphUntil(rdd) { (r, enqueue) => + val readsPipelined = r.dependencies.exists { + case _: PipelinedShuffleDependency[_, _, _] => true + case _: ShuffleDependency[_, _, _] => false // regular boundary: not within this stage + case narrowDep => + enqueue(narrowDep.rdd) + false + } + !readsPipelined + } } /** Invoke `.partitions` on the given RDD and all of its ancestors */ @@ -1992,9 +2041,13 @@ private[spark] class DAGScheduler( if (hasPipelined && rejectUnadmittablePipelinedGroup(jobId, finalRDD, partitions, listener)) { return } - var finalStage: ResultStage = null try { + // Reject group-level unsupported pipelined idioms (fan-out; spec S9) from the RDD graph, up + // front -- before any stage is created, so a rejection leaves no partial scheduler state. + // Inert for a job with no pipelined dependency. Inside this try so any incidental exception + // from the graph walk is handled by the same listener.jobFailed path as stage creation. + checkPipelinedGroupsSupportedInRDDGraph(finalRDD) // New stage creation may throw an exception if, for example, jobs are run on a // HadoopRDD whose underlying HDFS files have been deleted. finalStage = createResultStage(finalRDD, func, partitions, jobId, callSite) @@ -2027,6 +2080,14 @@ private[spark] class DAGScheduler( return } + case e: Exception with SparkThrowable if e.getCondition == "PIPELINED_SHUFFLE_UNSUPPORTED" => + // An up-front idiom rejection (checkPipelinedGroupsSupportedInRDDGraph / a producer-side + // check in createResultStage), not a stage-creation failure. Log it as such (the generic + // "Creating new stage failed" message below would be misleading). + logWarning(log"Rejecting job ${MDC(JOB_ID, jobId)}: unsupported pipelined-shuffle idiom", e) + listener.jobFailed(e) + return + case e: Exception => logWarning(log"Creating new stage failed due to exception - job: ${MDC(JOB_ID, jobId)}", e) listener.jobFailed(e) diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index 7e9c74cefd131..c6a68e341ff7f 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -41,7 +41,7 @@ import org.apache.spark.executor.ExecutorMetrics import org.apache.spark.internal.config import org.apache.spark.internal.config.{LEGACY_ABORT_STAGE_AFTER_KILL_TASKS, Tests} import org.apache.spark.network.shuffle.ExternalBlockStoreClient -import org.apache.spark.rdd.{DeterministicLevel, RDD} +import org.apache.spark.rdd.{DeterministicLevel, RDD, ReliableRDDCheckpointData} import org.apache.spark.resource.{ExecutorResourceRequests, ResourceProfile, ResourceProfileBuilder, TaskResourceProfile, TaskResourceRequests} import org.apache.spark.resource.ResourceUtils.{FPGA, GPU} import org.apache.spark.rpc.RpcTimeoutException @@ -7761,6 +7761,220 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assert(results === Map(0 -> 7, 1 -> 8), "the rerun must complete, re-committing its partitions") assertDataStructuresEmpty() } + + test("pipelined shuffle: buffered consumer TaskEnd events are still emitted when the job " + + "aborts") { + // A consumer's successful task completions are buffered (deferred) with their TaskEnd NOT yet + // posted while the producer runs. If the job is then torn down (here: the consumer's own task + // hits a FetchFailed, which aborts the group), those buffered successes must still emit their + // TaskEnd -- otherwise a listener that tracks active tasks (e.g. AppStatusListener, which only + // removes a stage once activeTasks hits 0) would leak the consumer stage as perpetually + // running. + // Mechanism: aborting the job tears down the still-running producer via + // cancelRunningIndependentStages, whose markStageAsFinished releases the consumer's deferral + // and + // drops its buffered events, posting their TaskEnd. This test guards that end-to-end invariant. + val endedTaskIds = new java.util.concurrent.ConcurrentHashMap[Long, Boolean]() + val recordingListener = new SparkListener { + override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = + endedTaskIds.put(taskEnd.taskInfo.taskId, true) + } + sc.addSparkListener(recordingListener) + try { + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = results.put(index, result) + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(consumerRdd, Array(0, 1), listener = failListener) + // Identify the co-scheduled consumer (ResultStage over consumerRdd) and its still-running + // pipelined producer by RDD identity rather than task-set order. + val consumerTaskSet = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + assert(scheduler.dependentStageMap.size === 1, + "the consumer must be co-scheduled with a running producer (a deferral must exist)") + + // Consumer partition 0 succeeds -> buffered (deferred), so no TaskEnd is posted yet. Give it + // a + // known taskId so we can assert its TaskEnd fires on teardown. + val bufferedTaskId = 7007L + runEvent(makeCompletionEvent(consumerTaskSet.tasks(0), Success, 42, + taskInfo = createFakeTaskInfoWithId(bufferedTaskId))) + sc.listenerBus.waitUntilEmpty() + assert(!endedTaskIds.containsKey(bufferedTaskId), + "the buffered consumer completion must not post its TaskEnd while deferred") + assert(scheduler.dependentStageMap.get( + scheduler.stageIdToStage(consumerTaskSet.stageId)).exists(_.delayedTaskCompletionEvents.nonEmpty), + "the consumer's successful completion must be buffered while its producer runs") + + // The consumer's OTHER task now hits a FetchFailed. For a pipelined group member this aborts + // the whole group (PR8 group-atomic failure) -- WITHOUT the producer ever finishing, so + // teardown goes through abortStage -> failJobAndIndependentStages -> + // cleanupStateForJobAndIndependentStages (NOT the producer-completion release path, which + // already posts TaskEnd correctly). The buffered success for partition 0 must still emit its + // TaskEnd on that cleanup path. + runEvent(makeCompletionEvent( + consumerTaskSet.tasks(1), + FetchFailed(makeBlockManagerId("hostA"), pipelinedDep.shuffleId, 0L, 0, 0, "ignored"), + null)) + scheduler.resubmitFailedStages() + assert(failure.get() != null, "the job must fail") + + // The buffered consumer success must still have produced a TaskEnd on teardown. + sc.listenerBus.waitUntilEmpty() + assert(endedTaskIds.containsKey(bufferedTaskId), + "a buffered consumer TaskEnd must still be emitted when the job aborts, to avoid leaking " + + "the stage as perpetually running in active-task listeners") + assertDataStructuresEmpty() + } finally { + sc.removeSparkListener(recordingListener) + } + } + + test("pipelined shuffle: a barrier producer stage is rejected") { + // A barrier stage exposes output only after a global sync, incompatible with incremental reads. + val producerRdd = new MyRDD(sc, 2, Nil).barrier().mapPartitions(iter => iter) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + assertPipelinedUnsupported(submitAndCaptureFailure(consumerRdd, Array(0, 1)), "barrier") + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a statically-indeterminate producer is rejected") { + // Indeterminate output's recovery is stage rollback-and-recompute, which a group never performs + // (moot under S6); reject rather than carry dead machinery. + val producerRdd = new MyRDD(sc, 2, Nil, indeterminate = true) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + assertPipelinedUnsupported(submitAndCaptureFailure(consumerRdd, Array(0, 1)), "indeterminate") + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a reliable RDD checkpoint in a PRODUCER's chain is rejected") { + // A reliable checkpoint writes a durable, lineage-truncated snapshot -> reintroduces cross-time + // reuse of a transient edge and needs a post-success recompute of the vanished input. Rejected + // by walking the producer's within-stage chain for a ReliableRDDCheckpointData. Keyed on + // checkpointData (not isCheckpointed): the write has not happened yet at group-creation time. + withTempDir { dir => + sc.setCheckpointDir(dir.getCanonicalPath) + val checkpointableRdd = new MyCheckpointRDD(sc, 2, Nil) + checkpointableRdd.checkpoint() // sets checkpointData to a ReliableRDDCheckpointData + assert(checkpointableRdd.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]])) + val pipelinedDep = new PipelinedShuffleDependency(checkpointableRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + assertPipelinedUnsupported( + submitAndCaptureFailure(consumerRdd, Array(0, 1)), "reliable RDD checkpoint") + assertDataStructuresEmpty() + } + } + + test("pipelined shuffle: a reliable RDD checkpoint in a CONSUMER's chain is rejected") { + // The rejection must cover a consumer member's chain too, not just the producer's: a consumer's + // transient input IS the pipelined shuffle, so a reliable checkpoint there would re-read the + // vanished stream on recompute. Pure all-PG shape (no regular shuffle -- v1 rejects those + // separately): producer --pipelined--> consumer(checkpointed, result). The consumer reads the + // pipelined shuffle and its within-stage chain carries the reliable checkpoint. + withTempDir { dir => + sc.setCheckpointDir(dir.getCanonicalPath) + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + // The consumer RDD reads the pipelined shuffle AND is reliably checkpointed. + val consumerRdd = new MyCheckpointRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + consumerRdd.checkpoint() + assert(consumerRdd.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]])) + assertPipelinedUnsupported( + submitAndCaptureFailure(consumerRdd, Array(0, 1)), "reliable RDD checkpoint") + assertDataStructuresEmpty() + } + } + + test("pipelined shuffle: a reliable checkpoint DOWNSTREAM in the consumer stage (not on the " + + "reading RDD) is still rejected") { + // Coverage for a checkpoint that is NOT on the RDD reading the pipelined shuffle, but on a + // narrow-dep child of it WITHIN the same consumer stage. Pure all-PG shape: + // producer --pipelined--> reads --(narrow)--> checkpointed (result). + // `reads` and `checkpointed` are one stage. Rooting the check at the pipelined-reading RDD and + // walking parents would MISS this (the checkpoint is downstream of the read); the check must + // instead recognize that `checkpointed`'s own within-stage chain reads a pipelined shuffle. + withTempDir { dir => + sc.setCheckpointDir(dir.getCanonicalPath) + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val readsRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + // A narrow-dep child of the reading RDD, in the SAME stage, that is reliably checkpointed, + // and is the result RDD (no regular shuffle after -- that would be a separately-rejected mix). + val checkpointedRdd = new MyCheckpointRDD(sc, 2, List(new OneToOneDependency(readsRdd))) + checkpointedRdd.checkpoint() + assert(checkpointedRdd.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]])) + assertPipelinedUnsupported( + submitAndCaptureFailure(checkpointedRdd, Array(0, 1)), "reliable RDD checkpoint") + assertDataStructuresEmpty() + } + } + + test("pipelined shuffle: a producer feeding more than one consumer (fan-out) is rejected") { + // 1:N fan-out is a supported model deferred to a later version; v1 rejects it. Fan-out is + // detected at the RDD level -- two DISTINCT RDDs listing the same pipelined shuffle as a + // dependency -- so it is expressible in a pure all-PG job without any regular shuffle: two + // consumer RDDs both read the same pipelined producer, unioned by a narrow dependency into the + // result. checkPipelinedGroupsSupportedInRDDGraph counts 2 distinct consumers for the shuffle. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerA = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + val consumerB = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + val union = new MyRDD(sc, 2, + List(new OneToOneDependency(consumerA), new OneToOneDependency(consumerB)), + tracker = mapOutputTracker) + assertPipelinedUnsupported(submitAndCaptureFailure(union, Array(0, 1)), "more than one consumer") + assertDataStructuresEmpty() + } + + test("regular shuffle idioms are NOT rejected (inertness of the pipelined fail-fast)") { + // The pipelined fail-fast checks (indeterminate producer, reliable-checkpoint-in-chain) must be + // inert for a job with NO pipelined dependency: a regular shuffle whose producer is BOTH + // indeterminate AND reliably checkpointed -- the two idioms most likely to false-fire -- must + // run exactly as before. (Actually exercise both, not just the flag.) + withTempDir { dir => + sc.setCheckpointDir(dir.getCanonicalPath) + val producerRdd = new MyCheckpointRDD(sc, 2, Nil, indeterminate = true) + producerRdd.checkpoint() + assert(producerRdd.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]])) + assert(producerRdd.outputDeterministicLevel == DeterministicLevel.INDETERMINATE) + val regularDep = new ShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(regularDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + // The job proceeds normally (producer stage submitted), i.e. NOT failed by a pipelined check. + assert(failure === null, "a regular shuffle must not be rejected by the pipelined fail-fast") + completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2) + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + } + + private def submitAndCaptureFailure(finalRdd: RDD[_], partitions: Array[Int]): Exception = { + val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = {} + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(finalRdd, partitions, listener = failListener) + failure.get() + } + + private def assertPipelinedUnsupported(failure: Exception, reasonSubstring: String): Unit = { + assert(failure != null, "the job must fail fast on the unsupported pipelined idiom") + val msg = failure.getMessage + assert(msg.contains("PIPELINED_SHUFFLE_UNSUPPORTED") || msg.contains("unsupported feature"), + s"expected a PIPELINED_SHUFFLE_UNSUPPORTED error, got: $msg") + assert(msg.contains(reasonSubstring), + s"expected reason to mention '$reasonSubstring', got: $msg") + } + } class DAGSchedulerAbortStageOffSuite extends DAGSchedulerSuite { 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 384cf9cdc5e03..4c7e8efc4eaba 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala @@ -3139,18 +3139,16 @@ class TaskSetManagerSuite // The TaskSetManager.executorLost "Resubmitted" loop re-enqueues an already-successful // ShuffleMapTask when its executor is lost and the map output looks gone -- for a decommission, // it looks gone whenever MapOutputTracker.getMapOutputLocation returns None. A pipelined - // shuffle - // is NEVER registered in the MapOutputTracker, so getMapOutputLocation is always None - // and - // this loop would resubmit the lone producer task -- the exact streaming-writer hang. This loop - // bypasses handleFailedTask, so the group-atomic abort would never see it. The guard + // shuffle is NEVER registered in the MapOutputTracker, so getMapOutputLocation is always None + // and this loop would resubmit the lone producer task -- the exact streaming-writer hang. This + // loop bypasses handleFailedTask, so the group-atomic abort would never see it. The guard // (!taskSet.isPipelined on maybeShuffleMapOutputLoss) must keep a pipelined set out of the // loop. // // Scenario: a 2-task pipelined producer, one task succeeds on the to-be-lost executor while the // other is still running (so the set is NOT a zombie and the loop is entered for a - // non-pipelined - // set). Decommission that executor. Assert NO Resubmitted is emitted for the completed task. + // non-pipelined set). Decommission that executor. Assert NO Resubmitted is emitted for the + // completed task. sc = new SparkContext("local", "test") sched = new FakeTaskScheduler(sc, ("execA", "host1"), ("execB", "host2")) sched.initialize(new FakeSchedulerBackend()) @@ -3205,8 +3203,7 @@ class TaskSetManagerSuite // Decommission execA (which ran the completed task 0) and lose it. For a NON-pipelined set this // would re-enqueue task 0 via Resubmitted (its output looks lost); for a pipelined set the - // guard - // must prevent that. + // guard must prevent that. manager.executorDecommission("execA") manager.executorLost("execA", "host1", ExecutorDecommission()) From c331c12496b9a1468814303811b1a2cc0f75e69b Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sun, 19 Jul 2026 07:26:41 +0000 Subject: [PATCH 17/39] [SPARK-XXXXX][CORE] Test: multi-producer consumer drops buffered successes when a producer fails Regression coverage for a review concern: a pipelined consumer with more than one pipelined producer (a join) must DROP its buffered successes when any producer fails, never replay them via a still-pending sibling. On the full stack a member failure is group-atomic (abortStage tears down all members and marks each producer failed), so releaseDeferredPipelinedConsumers sees producerFailed=true for every producer and drops -- the "sibling succeeds -> replay" trigger cannot occur. This test pins that behavior (a producer feeding two consumers would be fan-out and is rejected; two producers feeding one consumer is a supported all-PG join). Co-authored-by: Isaac --- .../spark/scheduler/DAGSchedulerSuite.scala | 48 +++++++++++++++++-- 1 file changed, 44 insertions(+), 4 deletions(-) diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index c6a68e341ff7f..5ff8d5c9715ad 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7129,6 +7129,45 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti s"a dropped fan-in consumer must not replay when a surviving producer succeeds; got $results") } + test("pipelined shuffle: a multi-producer consumer's buffered successes are dropped when one " + + "producer fails (no stale replay via a sibling)") { + // Isaac-review probe (group-atomic drop across MULTIPLE producers). Consumer C reads TWO + // pipelined producers P1 and P2 (a join; pure all-PG, no regular shuffle). C finishes early -> + // its successes are buffered against BOTH producers. If P1 then fails, the whole group must be + // torn down and C's buffered successes DROPPED -- they depended on P1's (now-invalid) output. + // The concern: releaseDeferredPipelinedConsumers evaluates producerFailed per finishing + // producer, so a naive impl could remove P1 (parents still has P2, no drop), then later see P2 + // "succeed" and REPLAY. Verify the shipped stack drops instead (group abort tears C down). + val rddP1 = new MyRDD(sc, 2, Nil) + val psdP1 = new PipelinedShuffleDependency(rddP1, new HashPartitioner(2)) + val rddP2 = new MyRDD(sc, 2, Nil) + val psdP2 = new PipelinedShuffleDependency(rddP2, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(psdP1, psdP2), tracker = mapOutputTracker) + val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = results.put(index, result) + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(consumerRdd, Array(0, 1), listener = failListener) + // P1, P2, and C are all co-scheduled (pure all-PG join, admitted up front). + assert(taskSets.size === 3, s"expected P1, P2, consumer co-scheduled, got ${taskSets.size}") + val tsP1 = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddP1).get + val tsC = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd).get + + // Consumer finishes early -> buffered against both P1 and P2. + complete(tsC, Seq((Success, 42), (Success, 43))) + assert(results.isEmpty, "consumer completions should be buffered while producers run") + assert(scheduler.dependentStageMap.keys.exists(_.rdd eq consumerRdd), + "consumer should be deferred against its two producers") + + // P1 fails. The group is torn down; C's buffered successes must be DROPPED, never replayed -- + // even though P2 has not (and now will not) complete. + failed(tsP1, "producer P1 blew up") + assert(failure.get() != null, "the job must fail when a pipelined producer fails") + assert(results.isEmpty, + "a multi-producer consumer's buffered successes must be dropped when any producer fails") + assertDataStructuresEmpty() + } test("pipelined shuffle: a deferred consumer task fires its TaskEnd exactly once (at replay)") { // A deferred CompletionEvent must have its side effects (task-end listener event, accumulator @@ -7807,8 +7846,8 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti sc.listenerBus.waitUntilEmpty() assert(!endedTaskIds.containsKey(bufferedTaskId), "the buffered consumer completion must not post its TaskEnd while deferred") - assert(scheduler.dependentStageMap.get( - scheduler.stageIdToStage(consumerTaskSet.stageId)).exists(_.delayedTaskCompletionEvents.nonEmpty), + assert(scheduler.dependentStageMap.get(scheduler.stageIdToStage(consumerTaskSet.stageId)) + .exists(_.delayedTaskCompletionEvents.nonEmpty), "the consumer's successful completion must be buffered while its producer runs") // The consumer's OTHER task now hits a FetchFailed. For a pipelined group member this aborts @@ -7906,7 +7945,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) val readsRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) // A narrow-dep child of the reading RDD, in the SAME stage, that is reliably checkpointed, - // and is the result RDD (no regular shuffle after -- that would be a separately-rejected mix). + // and is the result RDD (no regular shuffle after -- that would be separately rejected). val checkpointedRdd = new MyCheckpointRDD(sc, 2, List(new OneToOneDependency(readsRdd))) checkpointedRdd.checkpoint() assert(checkpointedRdd.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]])) @@ -7929,7 +7968,8 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti val union = new MyRDD(sc, 2, List(new OneToOneDependency(consumerA), new OneToOneDependency(consumerB)), tracker = mapOutputTracker) - assertPipelinedUnsupported(submitAndCaptureFailure(union, Array(0, 1)), "more than one consumer") + assertPipelinedUnsupported( + submitAndCaptureFailure(union, Array(0, 1)), "more than one consumer") assertDataStructuresEmpty() } From 86110e86d3b24bcc4f1407164939938b36e943b1 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sun, 19 Jul 2026 19:19:07 +0000 Subject: [PATCH 18/39] [SPARK-XXXXX][CORE] Review round 1: fix checkPipelinedProducerSupported doc on where group idioms are checked The scaladoc said group-level idioms (fan-out, mixed resource profiles, internal regular shuffle) are "checked separately where the group is co-scheduled". That is inaccurate: fan-out is rejected up front at job submission by checkPipelinedGroupsSupportedInRDDGraph (before any stage is created), not at co-scheduling time; and mixed resource profiles / an internal regular shuffle are structural invariants that do not arise for v1's single-profile all-pipelined shape and are not separately checked at all. Correct the doc to point at the right place and stop implying checks that do not exist. Co-authored-by: Isaac --- .../scala/org/apache/spark/scheduler/DAGScheduler.scala | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index bd5b6e7e3dc57..476350ad3a30c 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -724,9 +724,12 @@ private[spark] class DAGScheduler( * incremental reads. Rejecting here (before the stage is used) keeps a misuse from silently * mis-scheduling. Inert for a regular ShuffleDependency. * - * Group-level idioms (fan-out to more than one consumer, mixed resource profiles, a regular - * shuffle internal to a group) are checked separately where the group is co-scheduled, since they - * are properties of the group rather than a single producer stage. + * Group-level idioms are handled elsewhere, since they are properties of the group rather than a + * single producer stage: fan-out (a producer with more than one consumer) is rejected up front + * at job submission by `checkPipelinedGroupsSupportedInRDDGraph` (before any stage is created); + * mixed resource profiles and a regular shuffle internal to a group are structural invariants + * that do not arise for v1's single-profile, all-pipelined job shape and so are not checked + * (see `checkPipelinedGroupsSupportedInRDDGraph`). */ private def checkPipelinedProducerSupported(shuffleDep: ShuffleDependency[_, _, _]): Unit = { if (!shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) { From a109071b5df68e3ffb60b9d818ca51e75aa2ef7a Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 06:16:23 +0000 Subject: [PATCH 19/39] [SPARK-XXXXX][CORE] Test: a deferred consumer's TaskEnd fires inline, so a group abort cannot leak the stage Follow-up to the fine-grained deferral change (which lives on the scheduling PR): update the abort-path test that previously guarded the coarse model's re-emit-on-teardown behavior. Under the fine-grained model a deferred consumer's TaskEnd is posted in real time when the task finishes -- only the stage/job-completion bookkeeping is deferred -- so the "leak the stage as perpetually running" failure the coarse model had to patch at teardown cannot arise. This test exercises the group-atomic FetchFailed->abort path (which this PR introduces), so it belongs here rather than on the scheduling PR. The test now asserts the consumer's TaskEnd is delivered inline (before the producer's fate is known), and that on the subsequent group abort the deferred completion is dropped (its result is never applied) with no duplicate TaskEnd re-emitted. Co-authored-by: Isaac --- .../spark/scheduler/DAGSchedulerSuite.scala | 48 ++++++++----------- 1 file changed, 21 insertions(+), 27 deletions(-) diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index 5ff8d5c9715ad..72599237ed4a1 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7801,18 +7801,16 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } - test("pipelined shuffle: buffered consumer TaskEnd events are still emitted when the job " + - "aborts") { - // A consumer's successful task completions are buffered (deferred) with their TaskEnd NOT yet - // posted while the producer runs. If the job is then torn down (here: the consumer's own task - // hits a FetchFailed, which aborts the group), those buffered successes must still emit their - // TaskEnd -- otherwise a listener that tracks active tasks (e.g. AppStatusListener, which only - // removes a stage once activeTasks hits 0) would leak the consumer stage as perpetually - // running. - // Mechanism: aborting the job tears down the still-running producer via - // cancelRunningIndependentStages, whose markStageAsFinished releases the consumer's deferral - // and - // drops its buffered events, posting their TaskEnd. This test guards that end-to-end invariant. + test("pipelined shuffle: a deferred consumer TaskEnd fires inline, so an abort cannot leak the " + + "stage as running") { + // Fine-grained model (S5.1): a consumer's successful task posts its TaskEnd in real time as the + // task finishes -- only the stage/job-completion bookkeeping is deferred. So even if the job is + // later torn down before the producer finishes (here: the consumer's OTHER task hits a + // FetchFailed, which aborts the group), the earlier success's TaskEnd has ALREADY been emitted. + // A listener that tracks active tasks (e.g. AppStatusListener, which removes a stage once + // activeTasks hits 0) therefore never leaks the consumer stage as perpetually running -- the + // leak the coarse model had to patch by re-emitting buffered TaskEnds at teardown simply cannot + // arise here. val endedTaskIds = new java.util.concurrent.ConcurrentHashMap[Long, Boolean]() val recordingListener = new SparkListener { override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = @@ -7837,25 +7835,23 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assert(scheduler.dependentStageMap.size === 1, "the consumer must be co-scheduled with a running producer (a deferral must exist)") - // Consumer partition 0 succeeds -> buffered (deferred), so no TaskEnd is posted yet. Give it - // a - // known taskId so we can assert its TaskEnd fires on teardown. + // Consumer partition 0 succeeds. Its completion bookkeeping is deferred (no job result yet), + // but its TaskEnd is posted INLINE, right now. Give it a known taskId so we can assert that. val bufferedTaskId = 7007L runEvent(makeCompletionEvent(consumerTaskSet.tasks(0), Success, 42, taskInfo = createFakeTaskInfoWithId(bufferedTaskId))) sc.listenerBus.waitUntilEmpty() - assert(!endedTaskIds.containsKey(bufferedTaskId), - "the buffered consumer completion must not post its TaskEnd while deferred") + assert(endedTaskIds.containsKey(bufferedTaskId), + "the consumer completion's TaskEnd must be posted inline, as the task finishes") + assert(results.isEmpty, "the consumer's job result must still be deferred (producer running)") assert(scheduler.dependentStageMap.get(scheduler.stageIdToStage(consumerTaskSet.stageId)) .exists(_.delayedTaskCompletionEvents.nonEmpty), - "the consumer's successful completion must be buffered while its producer runs") + "the consumer's completion bookkeeping must be buffered while its producer runs") // The consumer's OTHER task now hits a FetchFailed. For a pipelined group member this aborts - // the whole group (PR8 group-atomic failure) -- WITHOUT the producer ever finishing, so - // teardown goes through abortStage -> failJobAndIndependentStages -> - // cleanupStateForJobAndIndependentStages (NOT the producer-completion release path, which - // already posts TaskEnd correctly). The buffered success for partition 0 must still emit its - // TaskEnd on that cleanup path. + // the whole group (group-atomic failure) WITHOUT the producer ever finishing. The already- + // deferred success is dropped (its result must not be applied), and no duplicate TaskEnd is + // emitted for it -- it was already delivered above. runEvent(makeCompletionEvent( consumerTaskSet.tasks(1), FetchFailed(makeBlockManagerId("hostA"), pipelinedDep.shuffleId, 0L, 0, 0, "ignored"), @@ -7863,11 +7859,9 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti scheduler.resubmitFailedStages() assert(failure.get() != null, "the job must fail") - // The buffered consumer success must still have produced a TaskEnd on teardown. sc.listenerBus.waitUntilEmpty() - assert(endedTaskIds.containsKey(bufferedTaskId), - "a buffered consumer TaskEnd must still be emitted when the job aborts, to avoid leaking " + - "the stage as perpetually running in active-task listeners") + assert(results.isEmpty, + "the deferred consumer success must be dropped on abort, not applied as a result") assertDataStructuresEmpty() } finally { sc.removeSparkListener(recordingListener) From 48db5e8dddd0bc1098024bd65a220e78e8fcff10 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 06:39:35 +0000 Subject: [PATCH 20/39] [SPARK-XXXXX][CORE] Pipelined fail-fast: match a typed exception, and reject mixed resource profiles (S9) Two review items on the pipelined-group fail-fast path: - Typed exception (was: brittle string match). handleJobSubmitted distinguished an up-front idiom rejection from a stage-creation failure by matching on e.getCondition == "PIPELINED_SHUFFLE_UNSUPPORTED" -- a rename or a wrapped cause would silently fall through to the generic "Creating new stage failed" branch and mis-report. Introduce a marker type PipelinedShuffleUnsupportedException (carrying the same error class, so the user-facing message is unchanged) and match on the TYPE. - Mixed-resource-profile fail-fast (spec S9). The gang slot check compares one demand against one profile's capacity (maxNumConcurrentTasks is per profile), so v1 requires a single-profile group. This was asserted in comments as a "structural invariant" but nothing enforced it. checkPipelinedGroupsSupportedInRDDGraph now collects the distinct explicit resource profiles across the group's RDDs in its existing single graph walk and rejects a group spanning more than one. Per-profile accounting remains a follow-up. Updated the two scaladocs that claimed mixed profiles "do not arise / are not checked". Also drops a stray internal milestone reference from a comment. Tests: a new test submits a producer/consumer group with two distinct resource profiles and asserts the up-front rejection; neutering the check makes it fail (non-vacuous). The existing idiom-rejection tests (barrier, indeterminate, checkpoint, fan-out) still pass through the typed-exception catch. Full DAGSchedulerSuite green (193). Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 64 +++++++++++++------ .../spark/scheduler/DAGSchedulerSuite.scala | 20 ++++++ 2 files changed, 65 insertions(+), 19 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index 476350ad3a30c..a5dc0311e76f0 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -710,10 +710,8 @@ private[spark] class DAGScheduler( stage } - private def pipelinedUnsupportedError(reason: String): SparkException = new SparkException( - errorClass = "PIPELINED_SHUFFLE_UNSUPPORTED", - messageParameters = scala.collection.immutable.Map("reason" -> reason), - cause = null) + private def pipelinedUnsupportedError(reason: String): PipelinedShuffleUnsupportedException = + new PipelinedShuffleUnsupportedException(reason) /** * Fail-fast on producer-side idioms a pipelined shuffle cannot support in v1 (spec S9), checked @@ -725,11 +723,11 @@ private[spark] class DAGScheduler( * mis-scheduling. Inert for a regular ShuffleDependency. * * Group-level idioms are handled elsewhere, since they are properties of the group rather than a - * single producer stage: fan-out (a producer with more than one consumer) is rejected up front - * at job submission by `checkPipelinedGroupsSupportedInRDDGraph` (before any stage is created); - * mixed resource profiles and a regular shuffle internal to a group are structural invariants - * that do not arise for v1's single-profile, all-pipelined job shape and so are not checked - * (see `checkPipelinedGroupsSupportedInRDDGraph`). + * single producer stage: fan-out (a producer with more than one consumer) and a mixed-resource- + * profile group are rejected up front at job submission by + * `checkPipelinedGroupsSupportedInRDDGraph` (before any stage is created). A regular shuffle + * internal to a group does not arise for v1's all-pipelined job shape (groups are split at + * regular-shuffle boundaries, S3) and so is not checked. */ private def checkPipelinedProducerSupported(shuffleDep: ShuffleDependency[_, _, _]): Unit = { if (!shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) { @@ -758,7 +756,7 @@ private[spark] class DAGScheduler( } // Push-based shuffle merge as the pipelined shuffle: exposes output only after a // post-completion finalize step, the opposite of incremental reads. A - // PipelinedShuffleDependency disables merge in its constructor (M1.6), so this is a defensive + // PipelinedShuffleDependency disables merge in its constructor, so this is a defensive // backstop against that being bypassed. if (shuffleDep.shuffleMergeEnabled) { throw pipelinedUnsupportedError("push-based shuffle merge as a pipelined shuffle") @@ -1153,14 +1151,15 @@ private[spark] class DAGScheduler( * shuffle. Checked here (not at stage creation) so a reject leaves no partial stage state, and * so BOTH the producer chain (rooted at pd.rdd) and each consumer chain (rooted at a consuming * RDD) are covered from the whole-graph view. + * - Mixed resource profiles across the group's RDDs. The gang slot check (S4) compares one + * demand against one profile's capacity, so v1 requires a group to be single-profile; more + * than one distinct explicit profile is rejected. Per-profile accounting is a follow-up. * - * (The other S9 group-level rows -- mixed resource profiles and a regular shuffle internal to a - * group -- are structural invariants that do not arise for the single-profile, - * prefix*->PG->suffix* shapes v1 targets: a group is single-profile by construction here, and - * groups are split at - * regular-shuffle boundaries (S3). The remaining producer-side idioms -- barrier, DRA, - * indeterminate, checksum, push-merge -- are rejected in checkPipelinedProducerSupported at stage - * creation, where a producer-only throw leaves no partial state.) + * (The remaining S9 group-level row -- a regular shuffle internal to a group -- is a structural + * invariant that does not arise for the prefix*->PG->suffix* shapes v1 targets: groups are split + * at regular-shuffle boundaries (S3). The producer-side idioms -- barrier, DRA, indeterminate, + * checksum, push-merge -- are rejected in checkPipelinedProducerSupported at stage creation, + * where a producer-only throw leaves no partial state.) */ private def checkPipelinedGroupsSupportedInRDDGraph(finalRDD: RDD[_]): Unit = { // Walk the whole RDD graph once, collecting for each pipelined shuffleId the distinct consumer @@ -1169,10 +1168,12 @@ private[spark] class DAGScheduler( val consumersByShuffleId = new HashMap[Int, HashSet[Int]] val producerRoots = new HashSet[RDD[_]] // RDDs that WRITE a pipelined shuffle val reliablyCheckpointed = new HashSet[RDD[_]] // RDDs with a reliable checkpoint pending + val resourceProfiles = new HashSet[ResourceProfile] // distinct RPs across the group's RDDs traverseRDDGraph(finalRDD) { (rdd, enqueue) => if (rdd.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]])) { reliablyCheckpointed += rdd } + Option(rdd.getResourceProfile()).foreach(resourceProfiles += _) rdd.dependencies.foreach { case pd: PipelinedShuffleDependency[_, _, _] => consumersByShuffleId.getOrElseUpdate(pd.shuffleId, new HashSet[Int]) += rdd.id @@ -1186,6 +1187,16 @@ private[spark] class DAGScheduler( throw pipelinedUnsupportedError( "a pipelined producer with more than one consumer (fan-out / branching)") } + // Mixed resource profiles within a group (spec S9). The gang slot check (S4) compares one + // demand against one profile's capacity (maxNumConcurrentTasks is defined per profile), so v1 + // requires a group to be single-profile and fails fast otherwise; per-profile accounting is a + // follow-up. An RDD with no explicit profile contributes nothing (it runs under the default), + // so only distinct EXPLICIT profiles are counted -- more than one means the group spans + // profiles. This is a real check, not just a documented assumption: nothing else enforces it. + if (resourceProfiles.size > 1) { + throw pipelinedUnsupportedError( + "members of a pipelined group with differing resource profiles (single-profile required)") + } // Reject a reliable checkpoint anywhere in a pipelined-group MEMBER's within-stage chain (spec // S9). Keyed on checkpointData being ReliableRDDCheckpointData, not isCheckpointed, since the // write has not happened yet. Cache / .persist() / local checkpoint are whole-partition and @@ -2083,10 +2094,11 @@ private[spark] class DAGScheduler( return } - case e: Exception with SparkThrowable if e.getCondition == "PIPELINED_SHUFFLE_UNSUPPORTED" => + case e: PipelinedShuffleUnsupportedException => // An up-front idiom rejection (checkPipelinedGroupsSupportedInRDDGraph / a producer-side // check in createResultStage), not a stage-creation failure. Log it as such (the generic - // "Creating new stage failed" message below would be misleading). + // "Creating new stage failed" message below would be misleading). Matched by TYPE, not by + // the error-condition string, so a rename or a wrapped cause cannot misroute it. logWarning(log"Rejecting job ${MDC(JOB_ID, jobId)}: unsupported pipelined-shuffle idiom", e) listener.jobFailed(e) return @@ -4703,6 +4715,20 @@ private[spark] object DAGScheduler { val RESUBMIT_TIMEOUT = 200 } +/** + * Thrown when a job uses a pipelined-shuffle idiom that v1 does not support (fan-out, a barrier / + * indeterminate / checksum-retry / push-merge producer, a reliable checkpoint in a member's chain, + * or a mixed-resource-profile group; spec S9). `handleJobSubmitted` matches on this TYPE to + * distinguish an up-front idiom rejection from an ordinary stage-creation failure, rather than on + * the error-condition string (which a rename or a wrapped cause would silently break). Carries the + * `PIPELINED_SHUFFLE_UNSUPPORTED` error class so the user-facing message is unchanged. + */ +private[scheduler] class PipelinedShuffleUnsupportedException(reason: String) + extends SparkException( + errorClass = "PIPELINED_SHUFFLE_UNSUPPORTED", + messageParameters = scala.collection.immutable.Map("reason" -> reason), + cause = null) + /** * A NOT thread-safe set that only keeps the last `capacity` elements added to it. */ diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index 72599237ed4a1..9ff0bbfdc69c4 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7967,6 +7967,26 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } + test("pipelined shuffle: a group whose members have differing resource profiles is rejected") { + // The gang slot check (S4) compares one demand against one profile's capacity, so v1 requires a + // group to be single-profile and fails fast otherwise (spec S9). Give the producer and consumer + // two distinct, non-default resource profiles: checkPipelinedGroupsSupportedInRDDGraph collects + // more than one distinct profile across the group's RDDs and rejects up front. + val rpA = new ResourceProfileBuilder() + .require(new ExecutorResourceRequests().cores(4)) + .require(new TaskResourceRequests().cpus(1)).build() + val rpB = new ResourceProfileBuilder() + .require(new ExecutorResourceRequests().cores(8)) + .require(new TaskResourceRequests().cpus(2)).build() + val producerRdd = new MyRDD(sc, 2, Nil).withResources(rpA) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + .withResources(rpB) + assertPipelinedUnsupported( + submitAndCaptureFailure(consumerRdd, Array(0, 1)), "differing resource profiles") + assertDataStructuresEmpty() + } + test("regular shuffle idioms are NOT rejected (inertness of the pipelined fail-fast)") { // The pipelined fail-fast checks (indeterminate producer, reliable-checkpoint-in-chain) must be // inert for a job with NO pipelined dependency: a regular shuffle whose producer is BOTH From d5f8686aae5cd76ac3d79b4b6c97506fdea50cea Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 17:34:41 +0000 Subject: [PATCH 21/39] [SPARK-XXXXX][CORE] Reject any non-default resource profile in a pipelined group; gate FetchFailed group check Round-3 review follow-ups: - Resource-profile admission gap (real). The mixed-profile check collected only distinct EXPLICIT profiles (getResourceProfile() is null when unset) and rejected on count > 1, while the slot check measures capacity against the DEFAULT profile. So a group running uniformly on one non-default profile, or a non-default member mixed with default members, passed the check yet was admitted against the wrong (default) capacity pool -- the exact partial-residency deadlock gang admission prevents. Reject if ANY member carries an explicit non-default profile: v1 requires the whole group on the default profile (spec S9); per-profile accounting remains a follow-up. Added tests for the two previously-missed shapes (uniform non-default; non-default + default), alongside the two-distinct case. - FetchFailed group-membership check now gated on ActiveJob.hasPipelinedDependency, matching the submitMissingTasks site, so a regular job's FetchFailed no longer walks the stage RDD graph. - Dropped an internal-review artifact ("Isaac-review probe") from a test comment. Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 50 ++++++++------ .../spark/scheduler/DAGSchedulerSuite.scala | 68 +++++++++++++------ 2 files changed, 78 insertions(+), 40 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index a5dc0311e76f0..6bdf6aea129d8 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -723,8 +723,8 @@ private[spark] class DAGScheduler( * mis-scheduling. Inert for a regular ShuffleDependency. * * Group-level idioms are handled elsewhere, since they are properties of the group rather than a - * single producer stage: fan-out (a producer with more than one consumer) and a mixed-resource- - * profile group are rejected up front at job submission by + * single producer stage: fan-out (a producer with more than one consumer) and a group with a + * non-default resource profile are rejected up front at job submission by * `checkPipelinedGroupsSupportedInRDDGraph` (before any stage is created). A regular shuffle * internal to a group does not arise for v1's all-pipelined job shape (groups are split at * regular-shuffle boundaries, S3) and so is not checked. @@ -1151,9 +1151,10 @@ private[spark] class DAGScheduler( * shuffle. Checked here (not at stage creation) so a reject leaves no partial stage state, and * so BOTH the producer chain (rooted at pd.rdd) and each consumer chain (rooted at a consuming * RDD) are covered from the whole-graph view. - * - Mixed resource profiles across the group's RDDs. The gang slot check (S4) compares one - * demand against one profile's capacity, so v1 requires a group to be single-profile; more - * than one distinct explicit profile is rejected. Per-profile accounting is a follow-up. + * - A non-default resource profile on any member. The gang slot check (S4) compares one demand + * against one profile's capacity and measures it against the default profile, so v1 requires + * the whole group to run on the default profile; any member with an explicit non-default + * profile is rejected. Per-profile accounting is a follow-up. * * (The remaining S9 group-level row -- a regular shuffle internal to a group -- is a structural * invariant that does not arise for the prefix*->PG->suffix* shapes v1 targets: groups are split @@ -1168,12 +1169,19 @@ private[spark] class DAGScheduler( val consumersByShuffleId = new HashMap[Int, HashSet[Int]] val producerRoots = new HashSet[RDD[_]] // RDDs that WRITE a pipelined shuffle val reliablyCheckpointed = new HashSet[RDD[_]] // RDDs with a reliable checkpoint pending - val resourceProfiles = new HashSet[ResourceProfile] // distinct RPs across the group's RDDs + var hasNonDefaultResourceProfile = false // any member RDD with a non-default RP traverseRDDGraph(finalRDD) { (rdd, enqueue) => if (rdd.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]])) { reliablyCheckpointed += rdd } - Option(rdd.getResourceProfile()).foreach(resourceProfiles += _) + // An RDD's EFFECTIVE profile is its explicit one, or the default when unset. The slot check + // measures capacity against the default profile (see rejectUnadmittablePipelinedGroup), so + // for v1 the whole group must run on the default profile; any explicit non-default profile on + // a member makes the group span profiles (against the default the rest use) and is rejected. + val rp = rdd.getResourceProfile() + if (rp != null && rp.id != ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID) { + hasNonDefaultResourceProfile = true + } rdd.dependencies.foreach { case pd: PipelinedShuffleDependency[_, _, _] => consumersByShuffleId.getOrElseUpdate(pd.shuffleId, new HashSet[Int]) += rdd.id @@ -1187,15 +1195,17 @@ private[spark] class DAGScheduler( throw pipelinedUnsupportedError( "a pipelined producer with more than one consumer (fan-out / branching)") } - // Mixed resource profiles within a group (spec S9). The gang slot check (S4) compares one - // demand against one profile's capacity (maxNumConcurrentTasks is defined per profile), so v1 - // requires a group to be single-profile and fails fast otherwise; per-profile accounting is a - // follow-up. An RDD with no explicit profile contributes nothing (it runs under the default), - // so only distinct EXPLICIT profiles are counted -- more than one means the group spans - // profiles. This is a real check, not just a documented assumption: nothing else enforces it. - if (resourceProfiles.size > 1) { + // Resource profile (spec S9). The gang slot check (S4) compares one demand against one + // profile's capacity (maxNumConcurrentTasks is defined per profile) and measures it against the + // DEFAULT profile, so v1 requires the whole group to run on the default profile and fails fast + // otherwise; per-profile accounting is a follow-up. Reject if ANY member carries an explicit + // non-default profile -- that both spans profiles (against the default the other members use) + // and would be admitted against the wrong (default) capacity pool. This is a real check, not + // just a documented assumption: nothing else enforces it. + if (hasNonDefaultResourceProfile) { throw pipelinedUnsupportedError( - "members of a pipelined group with differing resource profiles (single-profile required)") + "a pipelined group member with a non-default resource profile (v1 requires the default " + + "profile for the whole group)") } // Reject a reliable checkpoint anywhere in a pipelined-group MEMBER's within-stage chain (spec // S9). Keyed on checkpointData being ReliableRDDCheckpointData, not isCheckpointed, since the @@ -3376,8 +3386,10 @@ private[spark] class DAGScheduler( log"${MDC(STAGE_ATTEMPT_ID, task.stageAttemptId)} and there is a more recent attempt for " + log"that stage (attempt " + log"${MDC(NUM_ATTEMPT, failedStage.latestInfo.attemptNumber())}) running") - } else if (isPipelinedGroupMember(failedStage) || isPipelinedGroupMember(mapStage)) { - // Failure is group-atomic for a pipelined group. The base scheduler handles a + } else if (activeJobForStage(failedStage).flatMap(jobIdToActiveJob.get) + .exists(_.hasPipelinedDependency) && + (isPipelinedGroupMember(failedStage) || isPipelinedGroupMember(mapStage))) { + // Failure is group-atomic for a pipelined group (spec S6). The base scheduler handles a // FetchFailed by resubmitting just the map stage in isolation and recomputing serially, // but a transient pipelined shuffle cannot be re-read and its members are co-scheduled, so // a lone-stage resubmit is never valid and would deadlock the group. Abort the @@ -4718,8 +4730,8 @@ private[spark] object DAGScheduler { /** * Thrown when a job uses a pipelined-shuffle idiom that v1 does not support (fan-out, a barrier / * indeterminate / checksum-retry / push-merge producer, a reliable checkpoint in a member's chain, - * or a mixed-resource-profile group; spec S9). `handleJobSubmitted` matches on this TYPE to - * distinguish an up-front idiom rejection from an ordinary stage-creation failure, rather than on + * or a non-default resource profile on a member; spec S9). `handleJobSubmitted` matches on this + * TYPE to distinguish an up-front idiom rejection from an ordinary stage-creation failure, not on * the error-condition string (which a rename or a wrapped cause would silently break). Carries the * `PIPELINED_SHUFFLE_UNSUPPORTED` error class so the user-facing message is unchanged. */ diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index 9ff0bbfdc69c4..5a262f94a648a 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7131,13 +7131,13 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti test("pipelined shuffle: a multi-producer consumer's buffered successes are dropped when one " + "producer fails (no stale replay via a sibling)") { - // Isaac-review probe (group-atomic drop across MULTIPLE producers). Consumer C reads TWO - // pipelined producers P1 and P2 (a join; pure all-PG, no regular shuffle). C finishes early -> - // its successes are buffered against BOTH producers. If P1 then fails, the whole group must be - // torn down and C's buffered successes DROPPED -- they depended on P1's (now-invalid) output. - // The concern: releaseDeferredPipelinedConsumers evaluates producerFailed per finishing - // producer, so a naive impl could remove P1 (parents still has P2, no drop), then later see P2 - // "succeed" and REPLAY. Verify the shipped stack drops instead (group abort tears C down). + // Group-atomic drop across MULTIPLE producers. Consumer C reads TWO pipelined producers P1 and + // P2 (a join; pure all-PG, no regular shuffle). C finishes early -> its successes are buffered + // against BOTH producers. If P1 then fails, the whole group must be torn down and C's buffered + // successes DROPPED -- they depended on P1's (now-invalid) output. Note the release path + // evaluates producerFailed per finishing producer, so a naive impl could remove P1 + // (parents still has P2, no drop), then later see P2 "succeed" and REPLAY. This asserts the + // shipped stack drops instead: a member failure aborts the whole group, which tears C down. val rddP1 = new MyRDD(sc, 2, Nil) val psdP1 = new PipelinedShuffleDependency(rddP1, new HashPartitioner(2)) val rddP2 = new MyRDD(sc, 2, Nil) @@ -7967,23 +7967,49 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } - test("pipelined shuffle: a group whose members have differing resource profiles is rejected") { - // The gang slot check (S4) compares one demand against one profile's capacity, so v1 requires a - // group to be single-profile and fails fast otherwise (spec S9). Give the producer and consumer - // two distinct, non-default resource profiles: checkPipelinedGroupsSupportedInRDDGraph collects - // more than one distinct profile across the group's RDDs and rejects up front. - val rpA = new ResourceProfileBuilder() - .require(new ExecutorResourceRequests().cores(4)) - .require(new TaskResourceRequests().cpus(1)).build() - val rpB = new ResourceProfileBuilder() - .require(new ExecutorResourceRequests().cores(8)) - .require(new TaskResourceRequests().cpus(2)).build() - val producerRdd = new MyRDD(sc, 2, Nil).withResources(rpA) + // Resource-profile rejection (spec S9). The gang slot check measures capacity against the DEFAULT + // resource profile, so v1 requires the whole group to run on the default profile and rejects any + // member with an explicit non-default profile. The three shapes below must all be rejected; the + // uniform-non-default and non-default-plus-default cases in particular would each pass a + // "more than one DISTINCT profile" check yet still be admitted against the wrong (default) pool. + private def rpWithCores(cores: Int, cpus: Int): ResourceProfile = + new ResourceProfileBuilder() + .require(new ExecutorResourceRequests().cores(cores)) + .require(new TaskResourceRequests().cpus(cpus)).build() + + test("pipelined shuffle: a group with two distinct non-default resource profiles is rejected") { + val producerRdd = new MyRDD(sc, 2, Nil).withResources(rpWithCores(4, 1)) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + .withResources(rpWithCores(8, 2)) + assertPipelinedUnsupported( + submitAndCaptureFailure(consumerRdd, Array(0, 1)), "non-default resource profile") + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a group uniformly on one non-default resource profile is rejected") { + // Both members share ONE non-default profile -- a "distinct profile count" of 1, which a + // more-than-one-distinct-profile check would wrongly admit, then measure against the default + // profile's capacity (the wrong pool). Must be rejected: the whole group is off the default RP. + val rp = rpWithCores(4, 1) + val producerRdd = new MyRDD(sc, 2, Nil).withResources(rp) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + .withResources(rp) + assertPipelinedUnsupported( + submitAndCaptureFailure(consumerRdd, Array(0, 1)), "non-default resource profile") + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a group mixing a non-default profile with the default is rejected") { + // The producer carries an explicit non-default profile; the consumer is left on the default. + // The set of EXPLICIT profiles is again size 1, so a distinct-explicit-profile check would miss + // it, but the group genuinely spans the non-default and default profiles. Must be rejected. + val producerRdd = new MyRDD(sc, 2, Nil).withResources(rpWithCores(4, 1)) val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) - .withResources(rpB) assertPipelinedUnsupported( - submitAndCaptureFailure(consumerRdd, Array(0, 1)), "differing resource profiles") + submitAndCaptureFailure(consumerRdd, Array(0, 1)), "non-default resource profile") assertDataStructuresEmpty() } From 11f46a77fe638b0c12c37b8bbbc8c81ab72f4e22 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 21:21:47 +0000 Subject: [PATCH 22/39] [SPARK-XXXXX][CORE] Enforce the v1 job==group admission invariant; dedup a within-stage walk; warn when slot check is off Architecture-review follow-ups (no behavior change on the happy path): - Make the "job == one pipelined group, pre-admitted" invariant explicit and enforced. The submitStage co-schedule branch gang-schedules a group's members with no slot check, trusting that handleJobSubmitted already gang-admitted the whole job -- previously encoded only as a comment. Record ActiveJob.pipelinedGroupAdmitted on the admitted (or check-disabled) path and assert it at the co-schedule branch. This is a tripwire for a future version that relaxes job==group (multiple groups per job): it must move gang admission to a per-group check here and replace the assert, not delete it, or an unadmitted group could be gang-scheduled and deadlock. Verified the assert fires if admission is not recorded, and holds across all existing pipelined tests. - Collapse a duplicated within-stage walk: isPipelinedGroupMember's consumer walk was byte-identical to rddChainReadsPipelinedShuffle (differing only in the root RDD). Delegate to the latter as the single source of truth so a change to "what ends a within-stage chain" need not be made in two places. Behavior-preserving. - Warn (once) when spark.scheduler.pipelinedGroup.slotCheck.enabled=false, since disabling the only gang-admission deadlock check is safe only with out-of-band capacity reservation. Once-per-scheduler (event-loop-confined var) so it does not spam per submitted batch. Co-authored-by: Isaac --- .../apache/spark/scheduler/ActiveJob.scala | 14 +++++ .../apache/spark/scheduler/DAGScheduler.scala | 54 ++++++++++++------- 2 files changed, 50 insertions(+), 18 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/scheduler/ActiveJob.scala b/core/src/main/scala/org/apache/spark/scheduler/ActiveJob.scala index 7fa2b009fa018..91361d3262a55 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/ActiveJob.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/ActiveJob.scala @@ -71,4 +71,18 @@ private[spark] class ActiveJob( * short-circuit the group-membership graph walk for the common regular job at no cost. */ var hasPipelinedDependency: Boolean = false + + /** + * True once this job's pipelined group has passed up-front gang admission in + * `handleJobSubmitted` (or the slot check was disabled by config). This is a distinct fact from + * `hasPipelinedDependency` (which only says the job uses pipelining): the co-schedule path in + * `DAGScheduler.submitStage` gang-schedules a group's members with NO slot check, and asserts + * this flag so that trust in up-front admission is enforced rather than merely commented. + * + * v1 ONLY: the whole job is a single pipelined group, so admission is a job-level fact. If a + * later version allows multiple groups per job, this job-level flag must be replaced by + * per-group admission state, and the assert in `submitStage` replaced by that per-group gate -- + * NOT simply deleted, or an unadmitted group could be gang-scheduled and deadlock. + */ + var pipelinedGroupAdmitted: Boolean = false } diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index 6bdf6aea129d8..0e6f6009be4f4 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -187,6 +187,11 @@ private[spark] class DAGScheduler( delayedTaskCompletionEvents: ListBuffer[CompletionEvent] = new ListBuffer[CompletionEvent]) private[scheduler] val dependentStageMap = new HashMap[Stage, DependentStageInfo] + // Whether we have already logged that the pipelined-group slot check is disabled. Only the + // (single-threaded) event loop touches this, so a plain var is safe. Keeps the warning to once + // per scheduler rather than once per submitted batch job. + private var warnedPipelinedSlotCheckDisabled = false + private[scheduler] val activeJobs = new HashSet[ActiveJob] // Track all the jobs submitted by the same query execution, will clean up after @@ -1807,6 +1812,16 @@ private[spark] class DAGScheduler( return true } if (!sc.conf.get(config.PIPELINED_GROUP_SLOT_CHECK_ENABLED)) { + // The only deadlock-prevention check for gang admission is off. Legitimate only when the + // deployment admits capacity out-of-band (e.g. a slot reservation); otherwise a pipelined + // group that cannot co-fit will be gang-scheduled and can deadlock. Warn once so this is + // never a silent state. + if (!warnedPipelinedSlotCheckDisabled) { + warnedPipelinedSlotCheckDisabled = true + logWarning(log"${MDC(CONFIG, config.PIPELINED_GROUP_SLOT_CHECK_ENABLED.key)}=false: " + + log"pipelined-group gang admission is NOT checking free slots. This is safe only if " + + log"capacity is reserved out-of-band; otherwise a group that cannot co-fit may deadlock.") + } return false } val rp = sc.resourceProfileManager.defaultResourceProfile @@ -1870,24 +1885,12 @@ private[spark] class DAGScheduler( * member's `TaskSet.isPipelined` at submission and routing a member FetchFailed to a whole-group * abort. */ - private def isPipelinedGroupMember(stage: Stage): Boolean = { - if (isPipelinedProducer(stage)) { - return true - } - // Consumer check: does this stage read through a PipelinedShuffleDependency at one of its - // shuffle boundaries? Walk the stage's own RDD graph (descending narrow deps, stopping at every - // shuffle boundary) -- the same edges that define the stage -- and look for a pipelined one. - !traverseRDDGraphUntil(stage.rdd) { (rdd, enqueue) => - val hasPipelinedBoundary = rdd.dependencies.exists { - case _: PipelinedShuffleDependency[_, _, _] => true - case _: ShuffleDependency[_, _, _] => false // regular shuffle boundary: do not descend - case narrowDep => - enqueue(narrowDep.rdd) - false - } - !hasPipelinedBoundary // keep walking until a pipelined boundary is found - } - } + private def isPipelinedGroupMember(stage: Stage): Boolean = + // Producer side: it writes a pipelined shuffle. Consumer side: its within-stage chain reads + // one. rddChainReadsPipelinedShuffle is the single source of truth for that within-stage walk + // (descend narrow deps, stop at every shuffle boundary, look for a pipelined one) -- do not + // re-inline it; the consumer walk here and that method used to be byte-identical copies. + isPipelinedProducer(stage) || rddChainReadsPipelinedShuffle(stage.rdd) /** Finds the earliest-created active job that needs the stage */ // TODO: Probably should actually find among the active jobs that need this @@ -2125,6 +2128,11 @@ private[spark] class DAGScheduler( // Record whether this job uses a pipelined shuffle (computed above), so the per-submit // pipelined-group checks can short-circuit for a regular job without walking its RDD graph. job.hasPipelinedDependency = hasPipelined + // We only reach here if the up-front gang admission above (rejectUnadmittablePipelinedGroup) + // did NOT fail-and-return -- i.e. it passed or was disabled by config. Record that so the + // co-schedule path in submitStage can assert the group was admitted, rather than trusting a + // comment. (v1: job == one group, so admission is job-level; see ActiveJob.) + job.pipelinedGroupAdmitted = hasPipelined clearCacheLocs() logInfo( log"Got job ${MDC(JOB_ID, job.jobId)} (${MDC(CALL_SITE_SHORT_FORM, callSite.shortForm)}) " + @@ -2279,6 +2287,16 @@ private[spark] class DAGScheduler( pipelinedMissing.nonEmpty && pipelinedMissing.forall(runningStages.contains) if (regularMissing.isEmpty && allPipelinedParentsRunning) { + // v1 INVARIANT (job == one pipelined group): we co-schedule this group's members with + // NO slot check because the whole job was gang-admitted up front in + // handleJobSubmitted (rejectUnadmittablePipelinedGroup) before any member ran. Assert + // that link rather than trusting the comment. If a later version relaxes job==group + // (mixed DAGs / multiple groups per job), gang admission must move to a per-group + // "admit-on-ready" check HERE and REPLACE this assert -- do NOT just delete it, or an + // unadmitted group would be gang-scheduled and could deadlock. + assert(jobIdToActiveJob.get(jobId.get).exists(_.pipelinedGroupAdmitted), + s"Co-scheduling pipelined $stage whose job was not gang-admitted up front " + + "(the v1 job==group admission invariant is violated)") // The whole group's capacity was already admitted up front (handleJobSubmitted -> // rejectUnadmittablePipelinedGroup) before any member was submitted, so the group is // known to fit; just co-schedule this consumer with its running producer(s). No slot From 692da31dd639bf5c03873c3e13f95c74b74f6df5 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Tue, 21 Jul 2026 01:33:06 +0000 Subject: [PATCH 23/39] [SPARK-XXXXX][CORE] Make the job==group admission invariant fail-closed, not just an assert The co-schedule branch verified the up-front gang-admission invariant with a bare `assert`, which is elided under -da (JVM assertions off in production). In v1 the invariant always holds, so this is only a forward tripwire -- but if a later version relaxes job==group and forgets to move gang admission to a per-group check here, a production build would silently gang-schedule an unadmitted group and could deadlock. Replace the assert with a fail-closed abortStage (reusing the same terminal jobFailed path as the sibling no-active-job case), so the invariant is enforced even under -da: a violation becomes a clean, rerunnable job failure instead of a silent deadlock. Verified: the guard never fires in normal v1 operation (all pipelined tests green); neutering admission recording makes it abort the group rather than co-schedule. Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 26 ++++++++++++------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index 0e6f6009be4f4..951b487c4f57e 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -2286,17 +2286,23 @@ private[spark] class DAGScheduler( val allPipelinedParentsRunning = pipelinedMissing.nonEmpty && pipelinedMissing.forall(runningStages.contains) + // v1 INVARIANT (job == one pipelined group): we co-schedule this group's members with + // NO slot check because the whole job was gang-admitted up front in handleJobSubmitted + // (rejectUnadmittablePipelinedGroup) before any member ran. Verify that link rather + // than trusting the comment. In v1 this always holds (the flag is set on the same + // admitted path), so it is an invariant tripwire -- but make it FAIL-CLOSED (abort the + // group, not just assert): a bare `assert` is elided under -da, so if a later version + // relaxes job==group and forgets to move gang admission to a per-group "admit-on-ready" + // check HERE, a prod build would silently gang-schedule an unadmitted group and could + // deadlock. Aborting turns that into a clean, rerunnable job failure. (When that + // per-group check is added, it REPLACES this guard -- do not just delete it.) if (regularMissing.isEmpty && allPipelinedParentsRunning) { - // v1 INVARIANT (job == one pipelined group): we co-schedule this group's members with - // NO slot check because the whole job was gang-admitted up front in - // handleJobSubmitted (rejectUnadmittablePipelinedGroup) before any member ran. Assert - // that link rather than trusting the comment. If a later version relaxes job==group - // (mixed DAGs / multiple groups per job), gang admission must move to a per-group - // "admit-on-ready" check HERE and REPLACE this assert -- do NOT just delete it, or an - // unadmitted group would be gang-scheduled and could deadlock. - assert(jobIdToActiveJob.get(jobId.get).exists(_.pipelinedGroupAdmitted), - s"Co-scheduling pipelined $stage whose job was not gang-admitted up front " + - "(the v1 job==group admission invariant is violated)") + if (!jobIdToActiveJob.get(jobId.get).exists(_.pipelinedGroupAdmitted)) { + abortStage(stage, + s"Co-scheduling pipelined $stage whose job was not gang-admitted up front " + + "(the v1 job==group admission invariant is violated)", None) + return + } // The whole group's capacity was already admitted up front (handleJobSubmitted -> // rejectUnadmittablePipelinedGroup) before any member was submitted, so the group is // known to fit; just co-schedule this consumer with its running producer(s). No slot From 6945e2d0d64400c4424380ddeecf142b81685b48 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Tue, 21 Jul 2026 07:17:27 +0000 Subject: [PATCH 24/39] [SPARK-XXXXX][CORE] Remove internal jargon from pipelined-shuffle comments Strip internal shorthand that a reader of merged Apache Spark cannot resolve -- version/milestone tags, design-doc spec section numbers, "all-PG" abbreviation, and fine-grained/coarse model terminology -- from the pipelined-shuffle fail-fast comments and test comments, rewording to plain behavioral language. Comment and assertion-string text only; no behavioral change. Co-authored-by: Isaac --- .../apache/spark/scheduler/ActiveJob.scala | 12 +-- .../apache/spark/scheduler/DAGScheduler.scala | 97 ++++++++++--------- .../spark/scheduler/DAGSchedulerSuite.scala | 32 +++--- 3 files changed, 71 insertions(+), 70 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/scheduler/ActiveJob.scala b/core/src/main/scala/org/apache/spark/scheduler/ActiveJob.scala index 91361d3262a55..4491461ee1199 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/ActiveJob.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/ActiveJob.scala @@ -76,13 +76,13 @@ private[spark] class ActiveJob( * True once this job's pipelined group has passed up-front gang admission in * `handleJobSubmitted` (or the slot check was disabled by config). This is a distinct fact from * `hasPipelinedDependency` (which only says the job uses pipelining): the co-schedule path in - * `DAGScheduler.submitStage` gang-schedules a group's members with NO slot check, and asserts - * this flag so that trust in up-front admission is enforced rather than merely commented. + * `DAGScheduler.submitStage` gang-schedules a group's members with NO slot check, and checks this + * flag so that trust in up-front admission is enforced rather than merely commented. * - * v1 ONLY: the whole job is a single pipelined group, so admission is a job-level fact. If a - * later version allows multiple groups per job, this job-level flag must be replaced by - * per-group admission state, and the assert in `submitStage` replaced by that per-group gate -- - * NOT simply deleted, or an unadmitted group could be gang-scheduled and deadlock. + * The whole job is a single pipelined group, so admission is a job-level fact. If a future change + * allows multiple groups per job, this job-level flag must be replaced by per-group admission + * state, and the check in `submitStage` replaced by that per-group gate -- NOT simply deleted, or + * an unadmitted group could be gang-scheduled and deadlock. */ var pipelinedGroupAdmitted: Boolean = false } diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index 951b487c4f57e..7c1e257adc6b8 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -719,11 +719,11 @@ private[spark] class DAGScheduler( new PipelinedShuffleUnsupportedException(reason) /** - * Fail-fast on producer-side idioms a pipelined shuffle cannot support in v1 (spec S9), checked - * when the producer stage is created. A pipelined shuffle runs its producer and consumer stages - * concurrently over a transient, once-through stream that a group never recomputes in isolation - * (any failure aborts the whole group, S6), so mechanisms that recompute/roll back a single stage - * are moot, and features that expose output only after a global barrier are incompatible with + * Fail-fast on producer-side idioms a pipelined shuffle cannot support, checked when the producer + * stage is created. A pipelined shuffle runs its producer and consumer stages concurrently over a + * transient, once-through stream that a group never recomputes in isolation (any failure aborts + * the whole group), so mechanisms that recompute/roll back a single stage are moot, and features + * that expose output only after a global barrier are incompatible with * incremental reads. Rejecting here (before the stage is used) keeps a misuse from silently * mis-scheduling. Inert for a regular ShuffleDependency. * @@ -731,8 +731,8 @@ private[spark] class DAGScheduler( * single producer stage: fan-out (a producer with more than one consumer) and a group with a * non-default resource profile are rejected up front at job submission by * `checkPipelinedGroupsSupportedInRDDGraph` (before any stage is created). A regular shuffle - * internal to a group does not arise for v1's all-pipelined job shape (groups are split at - * regular-shuffle boundaries, S3) and so is not checked. + * internal to a group does not arise for the all-pipelined job shape (groups are split at + * regular-shuffle boundaries) and so is not checked. */ private def checkPipelinedProducerSupported(shuffleDep: ShuffleDependency[_, _, _]): Unit = { if (!shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) { @@ -749,7 +749,8 @@ private[spark] class DAGScheduler( throw pipelinedUnsupportedError("dynamic resource allocation with a pipelined shuffle") } // Statically-indeterminate producer: its recovery is stage rollback-and-recompute, which a - // group never performs (moot under S6); reject rather than carry dead machinery. + // group never performs (any failure aborts the whole group); reject rather than carry dead + // machinery. if (rdd.outputDeterministicLevel == DeterministicLevel.INDETERMINATE) { throw pipelinedUnsupportedError("a statically-indeterminate pipelined producer") } @@ -1138,34 +1139,34 @@ private[spark] class DAGScheduler( } /** - * Reject group-level idioms a pipelined group cannot support in v1 (spec S9), checked against the - * RDD graph BEFORE any stage is created -- so a rejection fails the job up front (via - * handleJobSubmitted's listener.jobFailed) without leaving partial scheduler state behind, - * exactly like the speculation check. + * Reject group-level idioms a pipelined group cannot support, checked against the RDD graph + * BEFORE any stage is created -- so a rejection fails the job up front (via handleJobSubmitted's + * listener.jobFailed) without leaving partial scheduler state behind, exactly like the + * speculation check. * * Inert for a job with no pipelined dependency. Throws PIPELINED_SHUFFLE_UNSUPPORTED on * violation. Enforces: - * - Fan-out: a pipelined producer feeding more than one consumer. 1:N is a supported model - * deferred to a later version (it needs multicast to N live readers); v1 rejects it. A + * - Fan-out: a pipelined producer feeding more than one consumer. 1:N is a supported model not + * yet built (it needs multicast to N live readers), so it is rejected for now. A * PipelinedShuffleDependency's producer is `dep.rdd`; a "consumer" is any RDD that lists that * dependency. More than one distinct consumer RDD for the same pipelined shuffle is fan-out. * - Reliable RDD checkpoint in a group member's within-stage chain (producer OR consumer side): * a reliable `checkpoint()` writes a durable, lineage-truncated snapshot, which both - * reintroduces cross-time reuse of a transient edge (S4) and requires a post-success recompute + * reintroduces cross-time reuse of a transient edge and requires a post-success recompute * of the member's transient input -- for a consumer, that input is the vanished pipelined * shuffle. Checked here (not at stage creation) so a reject leaves no partial stage state, and * so BOTH the producer chain (rooted at pd.rdd) and each consumer chain (rooted at a consuming * RDD) are covered from the whole-graph view. - * - A non-default resource profile on any member. The gang slot check (S4) compares one demand - * against one profile's capacity and measures it against the default profile, so v1 requires - * the whole group to run on the default profile; any member with an explicit non-default + * - A non-default resource profile on any member. The gang slot check compares one demand + * against one profile's capacity and measures it against the default profile, so the whole + * group is required to run on the default profile; any member with an explicit non-default * profile is rejected. Per-profile accounting is a follow-up. * - * (The remaining S9 group-level row -- a regular shuffle internal to a group -- is a structural - * invariant that does not arise for the prefix*->PG->suffix* shapes v1 targets: groups are split - * at regular-shuffle boundaries (S3). The producer-side idioms -- barrier, DRA, indeterminate, - * checksum, push-merge -- are rejected in checkPipelinedProducerSupported at stage creation, - * where a producer-only throw leaves no partial state.) + * (The remaining group-level case -- a regular shuffle internal to a group -- is a structural + * invariant that does not arise for the prefix -> pipelined-group -> suffix shapes targeted here: + * groups are split at regular-shuffle boundaries. The producer-side idioms -- barrier, DRA, + * indeterminate, checksum, push-merge -- are rejected in checkPipelinedProducerSupported at stage + * creation, where a producer-only throw leaves no partial state.) */ private def checkPipelinedGroupsSupportedInRDDGraph(finalRDD: RDD[_]): Unit = { // Walk the whole RDD graph once, collecting for each pipelined shuffleId the distinct consumer @@ -1181,8 +1182,8 @@ private[spark] class DAGScheduler( } // An RDD's EFFECTIVE profile is its explicit one, or the default when unset. The slot check // measures capacity against the default profile (see rejectUnadmittablePipelinedGroup), so - // for v1 the whole group must run on the default profile; any explicit non-default profile on - // a member makes the group span profiles (against the default the rest use) and is rejected. + // the whole group must run on the default profile; any explicit non-default profile on a + // member makes the group span profiles (against the default the rest use) and is rejected. val rp = rdd.getResourceProfile() if (rp != null && rp.id != ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID) { hasNonDefaultResourceProfile = true @@ -1200,20 +1201,20 @@ private[spark] class DAGScheduler( throw pipelinedUnsupportedError( "a pipelined producer with more than one consumer (fan-out / branching)") } - // Resource profile (spec S9). The gang slot check (S4) compares one demand against one - // profile's capacity (maxNumConcurrentTasks is defined per profile) and measures it against the - // DEFAULT profile, so v1 requires the whole group to run on the default profile and fails fast - // otherwise; per-profile accounting is a follow-up. Reject if ANY member carries an explicit - // non-default profile -- that both spans profiles (against the default the other members use) - // and would be admitted against the wrong (default) capacity pool. This is a real check, not - // just a documented assumption: nothing else enforces it. + // Resource profile. The gang slot check compares one demand against one profile's capacity + // (maxNumConcurrentTasks is defined per profile) and measures it against the DEFAULT profile, + // so the whole group is required to run on the default profile and fails fast otherwise; + // per-profile accounting is a follow-up. Reject if ANY member carries an explicit non-default + // profile -- that both spans profiles (against the default the other members use) and would be + // admitted against the wrong (default) capacity pool. This is a real check, not just a + // documented assumption: nothing else enforces it. if (hasNonDefaultResourceProfile) { throw pipelinedUnsupportedError( - "a pipelined group member with a non-default resource profile (v1 requires the default " + - "profile for the whole group)") + "a pipelined group member with a non-default resource profile (the whole group must run " + + "on the default profile)") } - // Reject a reliable checkpoint anywhere in a pipelined-group MEMBER's within-stage chain (spec - // S9). Keyed on checkpointData being ReliableRDDCheckpointData, not isCheckpointed, since the + // Reject a reliable checkpoint anywhere in a pipelined-group MEMBER's within-stage chain. + // Keyed on checkpointData being ReliableRDDCheckpointData, not isCheckpointed, since the // write has not happened yet. Cache / .persist() / local checkpoint are whole-partition and // ephemeral and are not rejected. A reliably-checkpointed RDD `cp` is inside a member stage // iff: @@ -2070,7 +2071,7 @@ private[spark] class DAGScheduler( } var finalStage: ResultStage = null try { - // Reject group-level unsupported pipelined idioms (fan-out; spec S9) from the RDD graph, up + // Reject group-level unsupported pipelined idioms (e.g. fan-out) from the RDD graph, up // front -- before any stage is created, so a rejection leaves no partial scheduler state. // Inert for a job with no pipelined dependency. Inside this try so any incidental exception // from the graph walk is handled by the same listener.jobFailed path as stage creation. @@ -2130,8 +2131,8 @@ private[spark] class DAGScheduler( job.hasPipelinedDependency = hasPipelined // We only reach here if the up-front gang admission above (rejectUnadmittablePipelinedGroup) // did NOT fail-and-return -- i.e. it passed or was disabled by config. Record that so the - // co-schedule path in submitStage can assert the group was admitted, rather than trusting a - // comment. (v1: job == one group, so admission is job-level; see ActiveJob.) + // co-schedule path in submitStage can verify the group was admitted, rather than trusting a + // comment. (The job is one group, so admission is job-level; see ActiveJob.) job.pipelinedGroupAdmitted = hasPipelined clearCacheLocs() logInfo( @@ -2286,12 +2287,12 @@ private[spark] class DAGScheduler( val allPipelinedParentsRunning = pipelinedMissing.nonEmpty && pipelinedMissing.forall(runningStages.contains) - // v1 INVARIANT (job == one pipelined group): we co-schedule this group's members with + // INVARIANT (job == one pipelined group): we co-schedule this group's members with // NO slot check because the whole job was gang-admitted up front in handleJobSubmitted // (rejectUnadmittablePipelinedGroup) before any member ran. Verify that link rather - // than trusting the comment. In v1 this always holds (the flag is set on the same + // than trusting the comment. This always holds today (the flag is set on the same // admitted path), so it is an invariant tripwire -- but make it FAIL-CLOSED (abort the - // group, not just assert): a bare `assert` is elided under -da, so if a later version + // group, not just assert): a bare `assert` is elided under -da, so if a future change // relaxes job==group and forgets to move gang admission to a per-group "admit-on-ready" // check HERE, a prod build would silently gang-schedule an unadmitted group and could // deadlock. Aborting turns that into a clean, rerunnable job failure. (When that @@ -2300,16 +2301,16 @@ private[spark] class DAGScheduler( if (!jobIdToActiveJob.get(jobId.get).exists(_.pipelinedGroupAdmitted)) { abortStage(stage, s"Co-scheduling pipelined $stage whose job was not gang-admitted up front " + - "(the v1 job==group admission invariant is violated)", None) + "(the job==group admission invariant is violated)", None) return } // The whole group's capacity was already admitted up front (handleJobSubmitted -> // rejectUnadmittablePipelinedGroup) before any member was submitted, so the group is // known to fit; just co-schedule this consumer with its running producer(s). No slot // check here -- that would re-measure capacity against a mid-flight snapshot and is - // unnecessary once admission is decided up front (v1 gang admission). Group-level + // unnecessary once admission is decided up front (gang admission). Group-level // idiom rejection (fan-out, internal regular shuffle) already happened at job - // submission (checkPipelinedGroupsSupportedInRDDGraph + the all-PG check). + // submission (checkPipelinedGroupsSupportedInRDDGraph + the all-pipelined check). logInfo(log"Submitting ${MDC(STAGE, stage)} concurrently with its running " + log"pipelined producer(s) ${MDC(MISSING_PARENT_STAGES, pipelinedMissing)}") // Record that this stage is co-scheduled with still-running pipelined producers, @@ -3413,7 +3414,7 @@ private[spark] class DAGScheduler( } else if (activeJobForStage(failedStage).flatMap(jobIdToActiveJob.get) .exists(_.hasPipelinedDependency) && (isPipelinedGroupMember(failedStage) || isPipelinedGroupMember(mapStage))) { - // Failure is group-atomic for a pipelined group (spec S6). The base scheduler handles a + // Failure is group-atomic for a pipelined group. The base scheduler handles a // FetchFailed by resubmitting just the map stage in isolation and recomputing serially, // but a transient pipelined shuffle cannot be re-read and its members are co-scheduled, so // a lone-stage resubmit is never valid and would deadlock the group. Abort the @@ -4752,9 +4753,9 @@ private[spark] object DAGScheduler { } /** - * Thrown when a job uses a pipelined-shuffle idiom that v1 does not support (fan-out, a barrier / + * Thrown when a job uses a pipelined-shuffle idiom that is not supported (fan-out, a barrier / * indeterminate / checksum-retry / push-merge producer, a reliable checkpoint in a member's chain, - * or a non-default resource profile on a member; spec S9). `handleJobSubmitted` matches on this + * or a non-default resource profile on a member). `handleJobSubmitted` matches on this * TYPE to distinguish an up-front idiom rejection from an ordinary stage-creation failure, not on * the error-condition string (which a rename or a wrapped cause would silently break). Carries the * `PIPELINED_SHUFFLE_UNSUPPORTED` error class so the user-facing message is unchanged. diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index 5a262f94a648a..adbff0e3315e0 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7132,8 +7132,8 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti test("pipelined shuffle: a multi-producer consumer's buffered successes are dropped when one " + "producer fails (no stale replay via a sibling)") { // Group-atomic drop across MULTIPLE producers. Consumer C reads TWO pipelined producers P1 and - // P2 (a join; pure all-PG, no regular shuffle). C finishes early -> its successes are buffered - // against BOTH producers. If P1 then fails, the whole group must be torn down and C's buffered + // P2 (a join; purely all-pipelined, no regular shuffle). C finishes early -> its successes are + // buffered against BOTH producers. If P1 then fails, the whole group must be torn down and C's // successes DROPPED -- they depended on P1's (now-invalid) output. Note the release path // evaluates producerFailed per finishing producer, so a naive impl could remove P1 // (parents still has P2, no drop), then later see P2 "succeed" and REPLAY. This asserts the @@ -7149,7 +7149,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti override def jobFailed(exception: Exception): Unit = failure.set(exception) } submit(consumerRdd, Array(0, 1), listener = failListener) - // P1, P2, and C are all co-scheduled (pure all-PG join, admitted up front). + // P1, P2, and C are all co-scheduled (purely all-pipelined join, admitted up front). assert(taskSets.size === 3, s"expected P1, P2, consumer co-scheduled, got ${taskSets.size}") val tsP1 = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddP1).get val tsC = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd).get @@ -7803,14 +7803,14 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti test("pipelined shuffle: a deferred consumer TaskEnd fires inline, so an abort cannot leak the " + "stage as running") { - // Fine-grained model (S5.1): a consumer's successful task posts its TaskEnd in real time as the + // A consumer's successful task posts its TaskEnd in real time as the // task finishes -- only the stage/job-completion bookkeeping is deferred. So even if the job is // later torn down before the producer finishes (here: the consumer's OTHER task hits a // FetchFailed, which aborts the group), the earlier success's TaskEnd has ALREADY been emitted. // A listener that tracks active tasks (e.g. AppStatusListener, which removes a stage once // activeTasks hits 0) therefore never leaks the consumer stage as perpetually running -- the - // leak the coarse model had to patch by re-emitting buffered TaskEnds at teardown simply cannot - // arise here. + // leak a defer-everything design would have to patch by re-emitting buffered TaskEnds at + // teardown simply cannot arise here. val endedTaskIds = new java.util.concurrent.ConcurrentHashMap[Long, Boolean]() val recordingListener = new SparkListener { override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = @@ -7879,7 +7879,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti test("pipelined shuffle: a statically-indeterminate producer is rejected") { // Indeterminate output's recovery is stage rollback-and-recompute, which a group never performs - // (moot under S6); reject rather than carry dead machinery. + // (so it never applies); reject rather than carry dead machinery. val producerRdd = new MyRDD(sc, 2, Nil, indeterminate = true) val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) @@ -7908,9 +7908,9 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti test("pipelined shuffle: a reliable RDD checkpoint in a CONSUMER's chain is rejected") { // The rejection must cover a consumer member's chain too, not just the producer's: a consumer's // transient input IS the pipelined shuffle, so a reliable checkpoint there would re-read the - // vanished stream on recompute. Pure all-PG shape (no regular shuffle -- v1 rejects those - // separately): producer --pipelined--> consumer(checkpointed, result). The consumer reads the - // pipelined shuffle and its within-stage chain carries the reliable checkpoint. + // vanished stream on recompute. Purely all-pipelined shape (no regular shuffle -- those are + // rejected separately): producer --pipelined--> consumer(checkpointed, result). The consumer + // reads the pipelined shuffle and its within-stage chain carries the reliable checkpoint. withTempDir { dir => sc.setCheckpointDir(dir.getCanonicalPath) val producerRdd = new MyRDD(sc, 2, Nil) @@ -7928,7 +7928,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti test("pipelined shuffle: a reliable checkpoint DOWNSTREAM in the consumer stage (not on the " + "reading RDD) is still rejected") { // Coverage for a checkpoint that is NOT on the RDD reading the pipelined shuffle, but on a - // narrow-dep child of it WITHIN the same consumer stage. Pure all-PG shape: + // narrow-dep child of it WITHIN the same consumer stage. Purely all-pipelined shape: // producer --pipelined--> reads --(narrow)--> checkpointed (result). // `reads` and `checkpointed` are one stage. Rooting the check at the pipelined-reading RDD and // walking parents would MISS this (the checkpoint is downstream of the read); the check must @@ -7950,9 +7950,9 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } test("pipelined shuffle: a producer feeding more than one consumer (fan-out) is rejected") { - // 1:N fan-out is a supported model deferred to a later version; v1 rejects it. Fan-out is + // 1:N fan-out is deferred to a later version and rejected up front here. Fan-out is // detected at the RDD level -- two DISTINCT RDDs listing the same pipelined shuffle as a - // dependency -- so it is expressible in a pure all-PG job without any regular shuffle: two + // dependency -- so it is expressible in an all-pipelined job without any regular shuffle: two // consumer RDDs both read the same pipelined producer, unioned by a narrow dependency into the // result. checkPipelinedGroupsSupportedInRDDGraph counts 2 distinct consumers for the shuffle. val producerRdd = new MyRDD(sc, 2, Nil) @@ -7967,9 +7967,9 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } - // Resource-profile rejection (spec S9). The gang slot check measures capacity against the DEFAULT - // resource profile, so v1 requires the whole group to run on the default profile and rejects any - // member with an explicit non-default profile. The three shapes below must all be rejected; the + // Resource-profile rejection. The gang slot check measures capacity against the DEFAULT + // resource profile, so the whole group must run on the default profile; any member with an + // explicit non-default profile is rejected. The three shapes below must all be rejected; the // uniform-non-default and non-default-plus-default cases in particular would each pass a // "more than one DISTINCT profile" check yet still be admitted against the wrong (default) pool. private def rpWithCores(cores: Int, cpus: Int): ResourceProfile = From a2c026759297f45a2944ee4f6adcbed1c500c6d3 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Thu, 23 Jul 2026 06:42:54 +0000 Subject: [PATCH 25/39] [SPARK-XXXXX][CORE] Supersede the pr3 resource-profile guard with the group-check version An earlier PR in this stack rejects a pipelined job whose member uses a non-default resource profile from rejectUnadmittablePipelinedGroup (a plain SparkException). This PR already performs the same rejection in checkPipelinedGroupsSupportedInRDDGraph, as a typed PipelinedShuffleUnsupported error alongside the fan-out and reliable-checkpoint group checks, and with broader coverage (uniform-non-default, mixed, and multi-profile shapes). Remove the now-redundant earlier check and its test so the resource-profile rejection lives in exactly one place with the typed error. Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 42 ------------- .../spark/scheduler/DAGSchedulerSuite.scala | 59 ------------------- 2 files changed, 101 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index 7c1e257adc6b8..05681e7f026a2 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -1770,48 +1770,6 @@ private[spark] class DAGScheduler( */ private def rejectUnadmittablePipelinedGroup( jobId: Int, finalRDD: RDD[_], partitions: Array[Int], listener: JobListener): Boolean = { - // Reject up-front (before any stage is created) two group members the admission model cannot - // support, walking the group's RDD graph once: - // - A barrier member. A barrier stage exposes its output only after a global sync, which - // contradicts a pipelined consumer reading the producer's output incrementally as it runs; - // and its failure recovery resubmits the stage, which the group's atomic completion cannot - // accommodate (a resubmitted producer would drop a co-scheduled consumer's buffered - // completions). Reject it here rather than let it be co-scheduled. - // - A member on a non-default resource profile. Admission below measures capacity and - // occupancy against the default profile, but each stage derives its profile from its RDDs - // (see createShuffleMapStage/createResultStage). A member carrying a non-default profile would be - // admitted against the default profile's free slots yet run in a different, often smaller - // pool -- and could then queue or deadlock there. - var offendingBarrier = false - var offendingRp: Option[ResourceProfile] = None - traverseRDDGraph(finalRDD) { (rdd, enqueue) => - if (rdd.isBarrier()) { - offendingBarrier = true - } - val rp = rdd.getResourceProfile() - if (offendingRp.isEmpty && rp != null && rp.id != DEFAULT_RESOURCE_PROFILE_ID) { - offendingRp = Some(rp) - } - rdd.dependencies.foreach(dep => enqueue(dep.rdd)) - } - if (offendingBarrier) { - logWarning(log"Rejecting job ${MDC(JOB_ID, jobId)}: a pipelined stage group contains a " + - log"barrier member") - listener.jobFailed(new SparkException( - "A pipelined shuffle job with a barrier stage in the pipelined group is not supported: a " + - "barrier stage exposes its output only after a global sync, which is incompatible with " + - "a pipelined consumer reading its output incrementally.")) - return true - } - if (offendingRp.nonEmpty) { - logWarning(log"Rejecting job ${MDC(JOB_ID, jobId)}: a pipelined stage group member uses a " + - log"non-default resource profile") - listener.jobFailed(new SparkException( - "A pipelined shuffle job with a member on a non-default resource profile is not " + - s"supported (resource profile id ${offendingRp.get.id}): the whole pipelined stage " + - "group must run on the default resource profile.")) - return true - } if (!sc.conf.get(config.PIPELINED_GROUP_SLOT_CHECK_ENABLED)) { // The only deadlock-prevention check for gang admission is off. Legitimate only when the // deployment admits capacity out-of-band (e.g. a slot reservation); otherwise a pipelined diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index adbff0e3315e0..3fe947e141962 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6586,65 +6586,6 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } } - test("pipelined shuffle: a member on a non-default resource profile is rejected up front") { - // Admission measures capacity/occupancy against the DEFAULT resource profile, but a stage - // derives its profile from its RDDs. A pipelined group member on a custom profile would be - // admitted against the default pool's free slots yet run in a different (often smaller) pool - // and could deadlock there. Reject such a job before any stage is created. - // Ensure the default profile exists (id 0) before building a custom one, so the custom profile - // gets a distinct, non-default id regardless of test-execution order (profile ids come from a - // process-wide counter, and the default profile occupies id 0). - sc.resourceProfileManager.defaultResourceProfile - val ereqs = new ExecutorResourceRequests().cores(4) - val treqs = new TaskResourceRequests().cpus(2) - val customRp = new ResourceProfileBuilder().require(ereqs).require(treqs).build() - val producerRdd = new MyRDD(sc, 2, Nil) - val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) - val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) - .withResources(customRp) - // Sanity: the consumer really carries a non-default profile (otherwise the test is vacuous). - assert(consumerRdd.getResourceProfile() != null && - consumerRdd.getResourceProfile().id != ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID, - s"test setup: expected a non-default profile, got id " + - s"${Option(consumerRdd.getResourceProfile()).map(_.id)}") - val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() - val failListener = new JobListener { - override def taskSucceeded(index: Int, result: Any): Unit = {} - override def jobFailed(exception: Exception): Unit = failure.set(exception) - } - submit(consumerRdd, Array(0, 1), listener = failListener) - assert(failure.get() != null, - "a pipelined job with a non-default-resource-profile member should be rejected") - assert(failure.get().getMessage.contains("non-default resource profile")) - assert(taskSets.isEmpty, "no stage should be created for a rejected job") - assertDataStructuresEmpty() - } - - test("pipelined shuffle: a barrier group member is rejected up front") { - // A barrier stage exposes its output only after a global sync, which is incompatible with a - // pipelined consumer reading the producer's output incrementally. It is also a recovery hazard: - // a barrier task failure fails the stage and RESUBMITS it (markStageAsFinished without - // willRetry), and if such a producer were co-scheduled, the resubmit would drop a fan-in - // consumer's buffered completions and the job could never complete. Reject any job whose - // pipelined group contains a barrier member before any stage is created, so that path is - // unreachable. - val producerRdd = new MyRDD(sc, 2, Nil).barrier().mapPartitions(iter => iter) - assert(producerRdd.isBarrier(), "test setup: the producer must be a barrier RDD") - val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) - val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) - val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() - val failListener = new JobListener { - override def taskSucceeded(index: Int, result: Any): Unit = {} - override def jobFailed(exception: Exception): Unit = failure.set(exception) - } - submit(consumerRdd, Array(0, 1), listener = failListener) - assert(failure.get() != null, - "a pipelined job with a barrier group member should be rejected") - assert(failure.get().getMessage.contains("barrier")) - assert(taskSets.isEmpty, "no stage should be created for a rejected job") - assertDataStructuresEmpty() - } - test("regular shuffle job with dynamic allocation enabled is NOT rejected (path is inert)") { // The dynamic-allocation fail-fast must apply only to jobs with a pipelined dependency; a plain // regular-shuffle job with dynamic allocation on runs normally. From 9f827340704c5ee0ca4b3d7142919a1e6f061d07 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Thu, 23 Jul 2026 17:22:57 +0000 Subject: [PATCH 26/39] [SPARK-XXXXX][CORE] Realign the group-abort deferral test with the coarse completion model The scheduling PR now buffers a pipelined consumer's whole completion event (coarse model) rather than running its TaskEnd inline and deferring only the finish bookkeeping. Update the group-abort test that asserted inline TaskEnd timing: - Before: "a deferred consumer TaskEnd fires inline, so an abort cannot leak the stage" -- asserted the buffered success's TaskEnd was emitted immediately, before the abort. - After: "a buffered consumer success is dropped on group abort, and its TaskEnd is flushed so the stage is not leaked as running" -- asserts no TaskEnd while the producer runs (the whole event is buffered), then that the FetchFailed group-abort's drop path flushes the buffered success's TaskEnd (so active-task-tracking listeners still see the task finish) while its result is dropped. The scenario (buffered success + sibling FetchFailed group-abort) and its no-leak guarantee are unchanged; only the event timing differs. Reverting the drop-path flush (events.foreach(postTaskEnd)) makes the test fail on the flush assertion. Co-authored-by: Isaac --- .../spark/scheduler/DAGSchedulerSuite.scala | 39 ++++++++++--------- 1 file changed, 21 insertions(+), 18 deletions(-) diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index 3fe947e141962..7f930edb686b9 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7742,16 +7742,16 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } - test("pipelined shuffle: a deferred consumer TaskEnd fires inline, so an abort cannot leak the " + - "stage as running") { - // A consumer's successful task posts its TaskEnd in real time as the - // task finishes -- only the stage/job-completion bookkeeping is deferred. So even if the job is - // later torn down before the producer finishes (here: the consumer's OTHER task hits a - // FetchFailed, which aborts the group), the earlier success's TaskEnd has ALREADY been emitted. - // A listener that tracks active tasks (e.g. AppStatusListener, which removes a stage once - // activeTasks hits 0) therefore never leaks the consumer stage as perpetually running -- the - // leak a defer-everything design would have to patch by re-emitting buffered TaskEnds at - // teardown simply cannot arise here. + test("pipelined shuffle: a buffered consumer success is dropped on group abort, and its " + + "TaskEnd is flushed so the stage is not leaked as running") { + // A consumer's successful task has its whole completion event buffered while its producer runs + // (coarse model) -- no TaskEnd yet. If the group is then torn down before the producer finishes + // (here: the consumer's OTHER task hits a FetchFailed, which aborts the group), the buffered + // success's result must NOT be applied, but its TaskEnd IS flushed on the drop path. A listener + // that tracks active tasks (e.g. AppStatusListener, which removes a stage once activeTasks hits + // 0) therefore does not leak the consumer stage as perpetually running: the success's TaskStart + // was delivered, and its matching TaskEnd arrives at teardown even though the result is + // dropped. val endedTaskIds = new java.util.concurrent.ConcurrentHashMap[Long, Boolean]() val recordingListener = new SparkListener { override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = @@ -7776,23 +7776,23 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assert(scheduler.dependentStageMap.size === 1, "the consumer must be co-scheduled with a running producer (a deferral must exist)") - // Consumer partition 0 succeeds. Its completion bookkeeping is deferred (no job result yet), - // but its TaskEnd is posted INLINE, right now. Give it a known taskId so we can assert that. + // Consumer partition 0 succeeds. Its whole completion event is buffered (no job result and no + // TaskEnd yet). Give it a known taskId so we can assert the drop path flushes it later. val bufferedTaskId = 7007L runEvent(makeCompletionEvent(consumerTaskSet.tasks(0), Success, 42, taskInfo = createFakeTaskInfoWithId(bufferedTaskId))) sc.listenerBus.waitUntilEmpty() - assert(endedTaskIds.containsKey(bufferedTaskId), - "the consumer completion's TaskEnd must be posted inline, as the task finishes") + assert(!endedTaskIds.containsKey(bufferedTaskId), + "the buffered consumer success must not emit its TaskEnd while the producer still runs") assert(results.isEmpty, "the consumer's job result must still be deferred (producer running)") assert(scheduler.dependentStageMap.get(scheduler.stageIdToStage(consumerTaskSet.stageId)) .exists(_.delayedTaskCompletionEvents.nonEmpty), - "the consumer's completion bookkeeping must be buffered while its producer runs") + "the consumer's completion event must be buffered while its producer runs") // The consumer's OTHER task now hits a FetchFailed. For a pipelined group member this aborts - // the whole group (group-atomic failure) WITHOUT the producer ever finishing. The already- - // deferred success is dropped (its result must not be applied), and no duplicate TaskEnd is - // emitted for it -- it was already delivered above. + // the whole group (group-atomic failure) WITHOUT the producer ever finishing. The buffered + // success is dropped (its result must not be applied), but its TaskEnd is flushed on the drop + // path so active-task-tracking listeners see the task finish. runEvent(makeCompletionEvent( consumerTaskSet.tasks(1), FetchFailed(makeBlockManagerId("hostA"), pipelinedDep.shuffleId, 0L, 0, 0, "ignored"), @@ -7803,6 +7803,9 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti sc.listenerBus.waitUntilEmpty() assert(results.isEmpty, "the deferred consumer success must be dropped on abort, not applied as a result") + assert(endedTaskIds.containsKey(bufferedTaskId), + "the buffered consumer success's TaskEnd must be flushed on the drop path, so the stage " + + "is not leaked as perpetually running") assertDataStructuresEmpty() } finally { sc.removeSparkListener(recordingListener) From 8fe4a72ca258fafad5202864979706e48182f8ff Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Thu, 23 Jul 2026 19:18:03 +0000 Subject: [PATCH 27/39] [SPARK-XXXXX][CORE] Gate the pipelined group-idiom check on hasPipelined so it cannot reject a regular job checkPipelinedGroupsSupportedInRDDGraph walks the whole RDD graph and rejects three group-level idioms: fan-out, a member on a non-default resource profile, and a reliable checkpoint in a member stage. The fan-out and checkpoint checks are keyed on a PipelinedShuffleDependency being present (consumersByShuffleId / producerRoots), so they are inert for a job with no pipelined dependency. The resource-profile check is NOT so keyed -- it fires for ANY RDD whose explicit profile is non-default. But handleJobSubmitted called checkPipelinedGroupsSupportedInRDDGraph unconditionally (the hasPipelined gates guarded only the mixed-shuffle and slot-admission checks). So an ordinary job with no pipelined dependency that merely attaches a non-default profile via RDD.withResources -- a GA stage-level-scheduling feature -- was rejected up front with PIPELINED_SHUFFLE_UNSUPPORTED "a pipelined group member with a non-default resource profile". That job runs fine on master; this was a regression wherever stage-level scheduling is used (K8s / YARN / Standalone with dynamic allocation). Fix: gate the whole checkPipelinedGroupsSupportedInRDDGraph call on hasPipelined, matching the two admission checks just above it. Every idiom it rejects concerns a pipelined group, so it must not run for a job that has no pipelined dependency. This changes behavior only for the previously mis-firing resource-profile branch (the other two were already inert); the three pipelined-group rejection tests still reject. Test: a regular job on a non-default resource profile (RDD.withResources, no pipelined dependency) must NOT be rejected and must run to completion. Reverting the gate makes it fail with PIPELINED_SHUFFLE_UNSUPPORTED. Full DAGSchedulerSuite green. Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 20 +++++++++++++------ .../spark/scheduler/DAGSchedulerSuite.scala | 20 +++++++++++++++++++ 2 files changed, 34 insertions(+), 6 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index 05681e7f026a2..6e0544bf28ad8 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -1144,8 +1144,10 @@ private[spark] class DAGScheduler( * listener.jobFailed) without leaving partial scheduler state behind, exactly like the * speculation check. * - * Inert for a job with no pipelined dependency. Throws PIPELINED_SHUFFLE_UNSUPPORTED on - * violation. Enforces: + * Call only for a job that has a pipelined dependency (handleJobSubmitted gates on + * hasPipelined): the resource-profile check below is not keyed on a pipelined dependency, so on a + * regular job it would reject an ordinary RDD.withResources(...) use. Throws + * PIPELINED_SHUFFLE_UNSUPPORTED on violation. Enforces: * - Fan-out: a pipelined producer feeding more than one consumer. 1:N is a supported model not * yet built (it needs multicast to N live readers), so it is rejected for now. A * PipelinedShuffleDependency's producer is `dep.rdd`; a "consumer" is any RDD that lists that @@ -2029,11 +2031,17 @@ private[spark] class DAGScheduler( } var finalStage: ResultStage = null try { - // Reject group-level unsupported pipelined idioms (e.g. fan-out) from the RDD graph, up - // front -- before any stage is created, so a rejection leaves no partial scheduler state. - // Inert for a job with no pipelined dependency. Inside this try so any incidental exception + // Reject group-level unsupported pipelined idioms (e.g. fan-out, a non-default resource + // profile, a reliable checkpoint in a member stage) from the RDD graph, up front -- before + // any stage is created, so a rejection leaves no partial scheduler state. Gated on + // hasPipelined: every idiom this checks concerns a pipelined group, so it must not run for a + // job with no pipelined dependency (the resource-profile check in particular is not keyed on + // a pipelined dependency and would otherwise reject an ordinary job that merely uses a + // non-default profile via RDD.withResources). Inside this try so any incidental exception // from the graph walk is handled by the same listener.jobFailed path as stage creation. - checkPipelinedGroupsSupportedInRDDGraph(finalRDD) + if (hasPipelined) { + checkPipelinedGroupsSupportedInRDDGraph(finalRDD) + } // New stage creation may throw an exception if, for example, jobs are run on a // HadoopRDD whose underlying HDFS files have been deleted. finalStage = createResultStage(finalRDD, func, partitions, jobId, callSite) diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index 7f930edb686b9..e1b3e608dcdb0 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7980,6 +7980,26 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } } + test("regular job on a non-default resource profile is NOT rejected (RP check is pipelined-only)") { + // The resource-profile rejection is not keyed on a pipelined dependency, so it must run ONLY + // for a job that has one (handleJobSubmitted gates checkPipelinedGroupsSupportedInRDDGraph on + // hasPipelined). A perfectly ordinary job that merely attaches a non-default profile via + // RDD.withResources -- a GA stage-level-scheduling feature -- has NO pipelined dependency and + // must run untouched. Without the gate the whole graph walk fires and rejects it with + // PIPELINED_SHUFFLE_UNSUPPORTED, a regression on plain withResources jobs. + val rdd = new MyRDD(sc, 2, Nil).withResources(rpWithCores(4, 2)) + // A non-default profile drives submitMissingTasks through addPySparkConfigsToProperties, which + // needs a non-null Properties; pass one (the default submit() overload leaves it null). + submit(rdd, Array(0, 1), properties = new Properties()) + assert(failure === null, + "a regular job with a non-default resource profile must not be rejected by the pipelined " + + "fail-fast") + assert(taskSets.nonEmpty, "the job must proceed to task submission, not be failed up front") + complete(taskSets(0), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + private def submitAndCaptureFailure(finalRdd: RDD[_], partitions: Array[Int]): Exception = { val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() val failListener = new JobListener { From 1d88366a044ab7753bef1aa9542cc55e213bcfb8 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Fri, 24 Jul 2026 00:36:40 +0000 Subject: [PATCH 28/39] [SPARK-XXXXX][CORE] Fix scalastyle line-length in the RP-inertness test name The test name "regular job on a non-default resource profile is NOT rejected (RP check is pipelined-only)" was 102 chars, over the 100-char scalastyle limit for test sources (core/Test/scalastyle). Shorten the parenthetical. Co-authored-by: Isaac --- .../scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index e1b3e608dcdb0..b8d804eb61ce8 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7980,7 +7980,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } } - test("regular job on a non-default resource profile is NOT rejected (RP check is pipelined-only)") { + test("regular job on a non-default resource profile is NOT rejected (RP check is pipelined)") { // The resource-profile rejection is not keyed on a pipelined dependency, so it must run ONLY // for a job that has one (handleJobSubmitted gates checkPipelinedGroupsSupportedInRDDGraph on // hasPipelined). A perfectly ordinary job that merely attaches a non-default profile via From ebcaa066f92fd65cc32c7c2c1c026b6e9c1f96e6 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Tue, 28 Jul 2026 05:40:35 +0000 Subject: [PATCH 29/39] [SPARK-XXXXX][CORE] Rework the fan-out admission-demand tests to assert the rejection this stack produces The diamond and wide-fan-out tests were written on the scheduling PR to prove the admission demand counts a reused pipelined producer once per shuffle id, not once per consumer edge. On that PR's base a fan-out group is admittable, so they asserted admission. This PR rejects fan-out outright (checkPipelinedGroupsSupportedInRDDGraph), so the admission assertions no longer hold once the two are stacked. Rework them to assert the outcome this stack actually produces while still proving the demand dedup. The up-front slot admission (rejectUnadmittablePipelinedGroup) runs before the fan-out idiom check, so at capacity 4 a correctly-deduped diamond (producer 2 + result 2 = 4) passes the slot check and is then rejected for fan-out (PIPELINED_SHUFFLE_UNSUPPORTED); a per-edge-counted group (6, or 8 for the wide fan-out) would be rejected earlier for CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT. The tests now assert the fan-out rejection AND that it is not a slot rejection, so a dedup regression flips the error class and fails them (revert-checked against per-edge counting). The separate "two distinct shuffles over one producer RDD" test is unchanged (it guards shuffle-id-vs-RDD dedup keying, not fan-out). Co-authored-by: Isaac --- .../spark/scheduler/DAGSchedulerSuite.scala | 52 +++++++++++-------- 1 file changed, 30 insertions(+), 22 deletions(-) diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index b8d804eb61ce8..05ba990e2d638 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6827,13 +6827,19 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } } - test("pipelined shuffle: a diamond reusing one producer counts its tasks once for admission") { + test("pipelined shuffle: admission counts a reused producer once (diamond rejected for " + + "fan-out, not for slots)") { // Fan-out/diamond: ONE PipelinedShuffleDependency (one producer, one shuffle id) is read by TWO - // consumers that are then narrow-joined into the result. Execution creates ONE producer stage - // (getOrCreateShuffleMapStage keys on shuffle id), so the group's real concurrent demand is - // producer(2) + result(2) = 4 -- NOT 6. Counting the producer once per consumer EDGE would make - // demand 6 and wrongly reject this group at capacity 4. Pinning capacity to exactly 4 admits - // the correctly-deduped group and would fail if the producer were double-counted. + // consumers that are then narrow-joined into the result. Fan-out is unsupported, so the group + // is ultimately rejected -- but WHICH rejection it gets proves the admission demand is deduped. + // The up-front slot admission (rejectUnadmittablePipelinedGroup) runs BEFORE the fan-out idiom + // check (checkPipelinedGroupsSupportedInRDDGraph). Execution creates ONE producer stage + // (getOrCreateShuffleMapStage keys on shuffle id), so the real concurrent demand is + // producer(2) + result(2) = 4. Counting the producer once per consumer EDGE would inflate it + // to 6. Pinning capacity to exactly 4: the deduped group PASSES the slot check and is then + // rejected for fan-out (PIPELINED_SHUFFLE_UNSUPPORTED); a double-counted group (6 > 4) would + // instead be rejected earlier for CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT. So a dedup regression + // flips the error class -- which this test detects. val producerRdd = new MyRDD(sc, 2, Nil) val myScheduler = scheduler.asInstanceOf[MyDAGScheduler] myScheduler.maxConcurrentTasksForTest = 4 @@ -6847,13 +6853,13 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti val resultRdd = new MyRDD( sc, 2, List(new OneToOneDependency(consumer1), new OneToOneDependency(consumer2)), tracker = mapOutputTracker) - submit(resultRdd, Array(0, 1)) - assert(taskSets.size === 2, - s"diamond group is producer + result (both consumers); got ${taskSets.size} task sets") - // Drain to completion: producer stage first, then the result stage. - completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2) - complete(taskSets(1), Seq((Success, 42), (Success, 43))) - assert(results === Map(0 -> 42, 1 -> 43)) + val failure = submitAndCaptureFailure(resultRdd, Array(0, 1)) + // Rejected for fan-out (demand fit within capacity 4), NOT for slots -- proving the reused + // producer was counted once. + assertPipelinedUnsupported(failure, "more than one consumer") + assert(!failure.getMessage.contains("CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT"), + s"a reused producer must be counted once (fit capacity 4), got a slot rejection: " + + s"${failure.getMessage}") assertDataStructuresEmpty() } finally { myScheduler.maxConcurrentTasksForTest = 1000 @@ -6861,10 +6867,13 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } } - test("pipelined shuffle: a wide fan-out reusing one producer counts its tasks once") { + test("pipelined shuffle: admission counts a reused producer once across a wide fan-out " + + "(rejected for fan-out, not for slots)") { // Wider fan-out: ONE producer feeds THREE consumers, all narrow-joined into the result. Real - // demand is still producer(2) + result(2) = 4; a per-edge count would be 2 + 3*2 = 8. Capacity - // 4 admits the deduped group and would reject the double-counted one. + // demand is still producer(2) + result(2) = 4; a per-edge count would be 2 + 3*2 = 8. At + // capacity 4 the deduped group passes the slot check and is rejected for fan-out; a + // double-counted group (8 > 4) would be rejected first for insufficient slots. See the diamond + // test above for the ordering rationale. val producerRdd = new MyRDD(sc, 2, Nil) val myScheduler = scheduler.asInstanceOf[MyDAGScheduler] myScheduler.maxConcurrentTasksForTest = 4 @@ -6875,12 +6884,11 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti new OneToOneDependency(new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)) }.toList val resultRdd = new MyRDD(sc, 2, consumers, tracker = mapOutputTracker) - submit(resultRdd, Array(0, 1)) - assert(taskSets.size === 2, - s"wide fan-out group is producer + result; got ${taskSets.size} task sets") - completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2) - complete(taskSets(1), Seq((Success, 42), (Success, 43))) - assert(results === Map(0 -> 42, 1 -> 43)) + val failure = submitAndCaptureFailure(resultRdd, Array(0, 1)) + assertPipelinedUnsupported(failure, "more than one consumer") + assert(!failure.getMessage.contains("CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT"), + s"a reused producer must be counted once (fit capacity 4), got a slot rejection: " + + s"${failure.getMessage}") assertDataStructuresEmpty() } finally { myScheduler.maxConcurrentTasksForTest = 1000 From 364a18eae7fffdd58bea419a235980c9b69e9fc9 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Wed, 22 Jul 2026 04:42:42 +0000 Subject: [PATCH 30/39] [SPARK-XXXXX][CORE] Register a pipelined shuffle with the streaming output tracker A pipelined shuffle streams records from a live producer stage to a co-scheduled consumer stage, so the consumer must look up its producer tasks' locations while both stages run. That routing lives in the StreamingShuffleOutputTracker, but createShuffleMapStage only registered the shuffle with the MapOutputTracker (which, for a transient pipelined shuffle, tracks an empty never-materialized output). As a result a consumer task querying for writer locations found the shuffle unregistered and failed. Register a PipelinedShuffleDependency's shuffle with the StreamingShuffleOutputTracker inside the same once-per-shuffleId guard as the MapOutputTracker registration. Inert for a regular shuffle (and when no streaming tracker is configured). Note: this is a latent gap in the co-scheduling layer independent of RTM; it belongs with the pipelined-scheduling change earlier in the stack and should be re-homed there (with a real-transport regression test) when the stack is next reorganized. Co-authored-by: Isaac --- .../org/apache/spark/scheduler/DAGScheduler.scala | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index 6e0544bf28ad8..d664a9fb3d96d 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -711,6 +711,21 @@ private[spark] class DAGScheduler( log"shuffle ${MDC(SHUFFLE_ID, shuffleDep.shuffleId)}") mapOutputTracker.registerShuffle(shuffleDep.shuffleId, rdd.partitions.length, shuffleDep.partitioner.numPartitions) + // A pipelined shuffle streams records from a live producer to a co-scheduled consumer, so the + // consumer must be able to look up its producer tasks' locations while both stages run. That + // routing lives in the StreamingShuffleOutputTracker (the MapOutputTracker registration above + // only tracks the empty, never-materialized durable output). Register here, inside the same + // once-per-shuffleId guard, so the tracker knows the producer/consumer fan-in before any + // consumer task asks for a writer location. Inert (tracker absent) for a regular shuffle. + if (shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) { + sc.env.streamingShuffleOutputTracker.foreach { tracker => + tracker.asInstanceOf[StreamingShuffleOutputTrackerMaster].registerShuffle( + shuffleDep.shuffleId, + rdd.partitions.length, + shuffleDep.partitioner.numPartitions, + jobId) + } + } } stage } From d6341d2163165e4166a04eac316147c57163a878 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Wed, 22 Jul 2026 04:43:06 +0000 Subject: [PATCH 31/39] [SPARK-XXXXX][SQL] Let ShuffleExchangeExec build a pipelined shuffle dependency Add a `pipelined` flag to ShuffleExchangeExec. When set, prepareShuffleDependency constructs a PipelinedShuffleDependency instead of a regular ShuffleDependency, so the DAGScheduler co-schedules the exchange's producer and consumer stages and routes the shuffle to the streaming shuffle manager. The flag defaults to false, so every existing shuffle is unchanged. Modeling this as a node field (rather than out-of-band metadata) keeps the pipelined/non-pipelined distinction part of the exchange's identity, so it participates in exchange-reuse and canonicalization -- correct for a transient shuffle that must not be reused as a durable one. Adding the field shifts ShuffleExchangeExec's positional arity, so update the pattern matches that destructure it positionally (EnsureRequirements, AQEUtils, InsertSortForLimitAndOffset and several test suites) to account for the new field; EnsureRequirements' repartition rewrite now uses copy() so the flag is preserved. Co-authored-by: Isaac --- .../InsertSortForLimitAndOffset.scala | 2 +- .../sql/execution/adaptive/AQEUtils.scala | 2 +- .../exchange/EnsureRequirements.scala | 6 +- .../exchange/ShuffleExchangeExec.scala | 42 +++++-- .../sql/DataFrameWindowFunctionsSuite.scala | 2 +- .../org/apache/spark/sql/DatasetSuite.scala | 2 +- .../spark/sql/execution/PlannerSuite.scala | 5 +- .../exchange/EnsureRequirementsSuite.scala | 110 +++++++++--------- .../sql/streaming/StreamingJoinSuite.scala | 4 +- 9 files changed, 97 insertions(+), 78 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffset.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffset.scala index aa29128cda7e0..88f0de2555f70 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffset.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffset.scala @@ -44,7 +44,7 @@ object InsertSortForLimitAndOffset extends Rule[SparkPlan] { _, // Should not match AQE shuffle stage because we only target un-submitted stages which // we can still rewrite the query plan. - s @ ShuffleExchangeExec(SinglePartition, child, _, _), + s @ ShuffleExchangeExec(SinglePartition, child, _, _, _), _) if child.logicalLink.isDefined => extractOrderingAndPropagateOrderingColumns(child) match { case Some((ordering, newChild)) => diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEUtils.scala index 578e0acd80525..cc3ca90739090 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEUtils.scala @@ -30,7 +30,7 @@ object AQEUtils { // Project/Filter/LocalSort/CollectMetrics. // Note: we only care about `HashPartitioning` as `EnsureRequirements` can only optimize out // user-specified repartition with `HashPartitioning`. - case ShuffleExchangeExec(h: HashPartitioning, _, shuffleOrigin, _) + case ShuffleExchangeExec(h: HashPartitioning, _, shuffleOrigin, _, _) if shuffleOrigin == REPARTITION_BY_COL || shuffleOrigin == REPARTITION_BY_NUM => val numPartitions = if (shuffleOrigin == REPARTITION_BY_NUM) { Some(h.numPartitions) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala index c632b3d841e61..8ca344e511946 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala @@ -272,8 +272,8 @@ case class EnsureRequirements( } child match { - case ShuffleExchangeExec(_, c, so, ps) => - ShuffleExchangeExec(newPartitioning, c, so, ps) + case s: ShuffleExchangeExec => + s.copy(outputPartitioning = newPartitioning) case gpe: GroupPartitionsExec => ShuffleExchangeExec(newPartitioning, gpe.child) case _ => ShuffleExchangeExec(newPartitioning, child) } @@ -896,7 +896,7 @@ case class EnsureRequirements( def apply(plan: SparkPlan): SparkPlan = { val newPlan = plan.transformUp { - case operator @ ShuffleExchangeExec(upper: HashPartitioning, child, shuffleOrigin, _) + case operator @ ShuffleExchangeExec(upper: HashPartitioning, child, shuffleOrigin, _, _) if optimizeOutRepartition && (shuffleOrigin == REPARTITION_BY_COL || shuffleOrigin == REPARTITION_BY_NUM) => def hasSemanticEqualPartitioning(partitioning: Partitioning): Boolean = { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala index dd829e697df61..92893a332ba08 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala @@ -191,7 +191,8 @@ case class ShuffleExchangeExec( override val outputPartitioning: Partitioning, child: SparkPlan, shuffleOrigin: ShuffleOrigin = ENSURE_REQUIREMENTS, - advisoryPartitionSize: Option[Long] = None) + advisoryPartitionSize: Option[Long] = None, + pipelined: Boolean = false) extends ShuffleExchangeLike { private lazy val writeMetrics = @@ -252,7 +253,8 @@ case class ShuffleExchangeExec( child.output, outputPartitioning, serializer, - writeMetrics) + writeMetrics, + pipelined) metrics("numPartitions").set(dep.partitioner.numPartitions) val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY) SQLMetrics.postDriverMetricUpdates( @@ -346,7 +348,8 @@ object ShuffleExchangeExec { outputAttributes: Seq[Attribute], newPartitioning: Partitioning, serializer: Serializer, - writeMetrics: Map[String, SQLMetric]) + writeMetrics: Map[String, SQLMetric], + pipelined: Boolean = false) : ShuffleDependency[Int, InternalRow, InternalRow] = { val part: Partitioner = newPartitioning match { case RoundRobinPartitioning(numPartitions) => new HashPartitioner(numPartitions) @@ -537,15 +540,30 @@ object ShuffleExchangeExec { } } val dependency = - new ShuffleDependency[Int, InternalRow, InternalRow]( - rddWithPartitionIds, - new PartitionIdPassthrough(part.numPartitions), - serializer, - shuffleWriterProcessor = createShuffleWriteProcessor(writeMetrics), - rowBasedChecksums = UnsafeRowChecksum.createUnsafeRowChecksums(checksumSize), - _checksumMismatchFullRetryEnabled = SQLConf.get.shuffleChecksumMismatchFullRetryEnabled, - checksumMismatchQueryLevelRollbackEnabled = - SQLConf.get.shuffleChecksumMismatchQueryLevelRollbackEnabled) + if (pipelined) { + // A pipelined shuffle is transient and incrementally readable: the DAGScheduler + // co-schedules its producer and consumer stages instead of materializing the shuffle + // first. The PipelinedShuffleDependency type is the entire opt-in -- routing to the + // streaming shuffle manager and pipelined-group co-scheduling both follow from it. The + // checksum-mismatch retry knobs are intentionally not carried over: a transient shuffle is + // never recomputed, so PipelinedShuffleDependency does not expose them (they stay off). + new PipelinedShuffleDependency[Int, InternalRow, InternalRow]( + rddWithPartitionIds, + new PartitionIdPassthrough(part.numPartitions), + serializer, + shuffleWriterProcessor = createShuffleWriteProcessor(writeMetrics), + rowBasedChecksums = UnsafeRowChecksum.createUnsafeRowChecksums(checksumSize)) + } else { + new ShuffleDependency[Int, InternalRow, InternalRow]( + rddWithPartitionIds, + new PartitionIdPassthrough(part.numPartitions), + serializer, + shuffleWriterProcessor = createShuffleWriteProcessor(writeMetrics), + rowBasedChecksums = UnsafeRowChecksum.createUnsafeRowChecksums(checksumSize), + _checksumMismatchFullRetryEnabled = SQLConf.get.shuffleChecksumMismatchFullRetryEnabled, + checksumMismatchQueryLevelRollbackEnabled = + SQLConf.get.shuffleChecksumMismatchQueryLevelRollbackEnabled) + } dependency } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala index c1d1f13915a6e..136dc57fdfcfa 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala @@ -1360,7 +1360,7 @@ class DataFrameWindowFunctionsSuite extends SharedSparkSession def isShuffleExecByRequirement( plan: ShuffleExchangeExec, desiredClusterColumns: Seq[String]): Boolean = plan match { - case ShuffleExchangeExec(op: HashPartitioning, _, ENSURE_REQUIREMENTS, _) => + case ShuffleExchangeExec(op: HashPartitioning, _, ENSURE_REQUIREMENTS, _, _) => partitionExpressionsColumns(op.expressions) === desiredClusterColumns case _ => false } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala index 3b28cae31a134..096eee137b320 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala @@ -1942,7 +1942,7 @@ class DatasetSuite extends SharedSparkSession val agg = cp.groupBy($"id" % 2).agg(count($"id")) agg.queryExecution.executedPlan.collectFirst { - case ShuffleExchangeExec(_, _: RDDScanExec, _, _) => + case ShuffleExchangeExec(_, _: RDDScanExec, _, _, _) => case BroadcastExchangeExec(_, _: RDDScanExec) => }.foreach { _ => fail( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala index 71192e92af0d0..41b74c6d4ff21 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala @@ -733,10 +733,11 @@ class PlannerSuite extends SharedSparkSession with AdaptiveSparkPlanHelper { outputPlan match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, - ShuffleExchangeExec(HashPartitioning(leftPartitioningExpressions, _), _, _, _), _), + ShuffleExchangeExec(HashPartitioning(leftPartitioningExpressions, _), + _, _, _, _), _), SortExec(_, _, ShuffleExchangeExec(HashPartitioning(rightPartitioningExpressions, _), - _, _, _), _), _) => + _, _, _, _), _), _) => assert(leftKeys === smjExec.leftKeys) assert(rightKeys === smjExec.rightKeys) assert(leftKeys === leftPartitioningExpressions) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala index 17d00ec055e07..95085b71bc251 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala @@ -60,7 +60,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(smjExec1) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: PartitioningCollection, _, _), _), - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB)) assert(rightKeys === Seq(exprB, exprA)) case other => fail(other.toString) @@ -71,7 +71,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprB :: exprA :: Nil, Inner, None, plan2, plan1) EnsureRequirements.apply(smjExec2) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, _: PartitioningCollection, _, _), _), _) => assert(leftKeys === Seq(exprB, exprA)) assert(rightKeys === Seq(exprA, exprB)) @@ -84,7 +84,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprD :: exprC :: Nil, exprB :: exprA :: Nil, Inner, None, plan1, plan1) EnsureRequirements.apply(smjExec3) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, _: PartitioningCollection, _, _), _), _) => assert(leftKeys === Seq(exprC, exprD)) assert(rightKeys === Seq(exprA, exprB)) @@ -126,8 +126,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { ) EnsureRequirements.apply(smjExec2) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprC, exprB, exprD)) assert(rightKeys === Seq(exprD, exprA, exprC)) case other => fail(other.toString) @@ -145,7 +145,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprB :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec1) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), _) => assert(leftKeys === Seq(exprB, exprA)) assert(rightKeys === Seq(exprB, exprC)) @@ -159,7 +159,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprB :: Nil, Inner, None, plan1, plan3) EnsureRequirements.apply(smjExec2) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, _: PartitioningCollection, _, _), _), _) => assert(leftKeys === Seq(exprB, exprA)) assert(rightKeys === Seq(exprB, exprC)) @@ -173,7 +173,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(smjExec3) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: PartitioningCollection, _, _), _), - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprB, exprC)) assert(rightKeys === Seq(exprB, exprA)) case other => fail(other.toString) @@ -319,7 +319,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB)) assert(rightKeys === Seq(exprC, exprD)) assert(p.expressions == Seq(exprC)) @@ -335,7 +335,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB)) assert(rightKeys === Seq(exprC, exprD)) @@ -353,7 +353,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB)) assert(rightKeys === Seq(exprC, exprD)) assert(p.expressions == Seq(exprC)) @@ -372,7 +372,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB, exprB)) assert(rightKeys === Seq(exprA, exprC, exprC)) assert(p.expressions == Seq(exprA, exprC, exprA)) @@ -388,7 +388,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB, exprB)) assert(rightKeys === Seq(exprA, exprC, exprD)) assert(p.expressions == Seq(exprA, exprC, exprA)) @@ -444,8 +444,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.numPartitions == 5) assert(right.numPartitions == 5) case other => fail(other.toString) @@ -462,7 +462,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, DummySparkPlan(_, _, left: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.numPartitions == 10) assert(right.numPartitions == 10) assert(right.expressions == Seq(exprC, exprD)) @@ -481,8 +481,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.numPartitions == 5) assert(left.expressions == Seq(exprA, exprB)) assert(right.numPartitions == 5) @@ -492,7 +492,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, DummySparkPlan(_, _, left: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.numPartitions == 1) assert(right.numPartitions == 1) assert(right.expressions == Seq(exprC)) @@ -510,8 +510,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.numPartitions == conf.numShufflePartitions) assert(left.expressions == Seq(exprA, exprB)) assert(right.numPartitions == conf.numShufflePartitions) @@ -529,7 +529,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, DummySparkPlan(_, _, left: PartitioningCollection, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.numPartitions == 10) assert(right.numPartitions == 10) assert(right.expressions == Seq(exprA)) @@ -545,7 +545,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { Inner, None, plan1, plan2) applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, right: PartitioningCollection, _, _), _), _) => assert(left.numPartitions == 20) assert(left.expressions == Seq(exprC)) @@ -587,7 +587,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, DummySparkPlan(_, _, left: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB)) assert(right.expressions === Seq(exprC, exprD)) assert(left.numPartitions == 6) @@ -606,7 +606,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, right: HashPartitioning, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB)) assert(right.expressions === Seq(exprC, exprD)) @@ -624,7 +624,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, right: HashPartitioning, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB)) assert(right.expressions === Seq(exprC, exprD)) @@ -644,7 +644,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, right: HashPartitioning, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB)) assert(right.expressions === Seq(exprC, exprD)) @@ -666,8 +666,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB)) assert(right.expressions === Seq(exprC, exprD)) assert(left.numPartitions == conf.numShufflePartitions) @@ -691,7 +691,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, left: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB)) assert(rightKeys === Seq(exprC, exprD)) assert(left.numPartitions == 9) @@ -714,8 +714,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { var smjExec = SortMergeJoinExec(exprA :: Nil, exprC :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA)) assert(rightKeys === Seq(exprC)) assert(left.numPartitions == 20) @@ -733,7 +733,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA)) assert(rightKeys === Seq(exprC)) assert(right.numPartitions == 10) @@ -765,8 +765,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { } else { EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA)) assert(rightKeys === Seq(exprC)) assert(left.numPartitions == 5) @@ -875,8 +875,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB, exprB)) assert(right.expressions === Seq(exprA, exprC, exprC)) case other => fail(other.toString) @@ -951,8 +951,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: exprC :: Nil, exprA :: exprB :: exprC :: Nil, Inner, None, plan1, plan2) applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB, exprC)) assert(right.expressions === Seq(exprA, exprB, exprC)) case other => fail(other.toString) @@ -969,8 +969,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB, exprB)) assert(right.expressions === Seq(exprA, exprC, exprC)) case other => fail(other.toString) @@ -987,8 +987,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB, exprB)) assert(right.expressions === Seq(exprA, exprC, exprC)) case other => fail(other.toString) @@ -1008,8 +1008,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB, exprB)) assert(right.expressions === Seq(exprA, exprC, exprC)) case other => fail(other.toString) @@ -1130,7 +1130,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { case ShuffledHashJoinExec(_, _, _, _, _, DummySparkPlan(_, _, left: KeyedPartitioning, _, _), ShuffleExchangeExec(KeyedPartitioning(attrs, pks, _, _), - DummySparkPlan(_, _, SinglePartition, _, _), _, _), _) => + DummySparkPlan(_, _, SinglePartition, _, _), _, _, _), _) => assert(left.expressions == a1 :: Nil) assert(attrs == a1 :: Nil) assert(partitionKeys == pks.map(_.row)) @@ -1234,8 +1234,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(join) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(p1: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(p2: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(p1: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(p2: HashPartitioning, _, _, _, _), _), _) => // Both sides should be shuffled to default partitions assert(p1.numPartitions == 10) assert(p2.numPartitions == 10) @@ -1259,8 +1259,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(join) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), _) => // Both sides shuffled due to key mismatch case other => fail(s"Expected shuffles on both sides, but got: $other") } @@ -1278,7 +1278,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(join) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), SortExec(_, _, _: DummySparkPlan, _), _) => // Left side shuffled, right side kept as-is case other => fail(s"Expected shuffle on the left side, but got: $other") @@ -1296,8 +1296,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(join) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), _) => // Both sides shuffled due to canCreatePartitioning = false case other => fail(s"Expected shuffles on both sides, but got: $other") } @@ -1371,8 +1371,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(join) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(p1: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(p2: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(p1: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(p2: HashPartitioning, _, _, _, _), _), _) => // Both sides should be shuffled because partition keys not in join keys assert(p1.numPartitions == 10) assert(p2.numPartitions == 10) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala index c46f0076721b9..c0983f338abe5 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala @@ -677,8 +677,8 @@ abstract class StreamingInnerJoinBase extends StreamingJoinSuite { assert(query.lastExecution.executedPlan.collect { case j @ StreamingSymmetricHashJoinExec(_, _, _, _, _, _, _, _, _, - ShuffleExchangeExec(opA: HashPartitioning, _, _, _), - ShuffleExchangeExec(opB: HashPartitioning, _, _, _), _) + ShuffleExchangeExec(opA: HashPartitioning, _, _, _, _), + ShuffleExchangeExec(opB: HashPartitioning, _, _, _, _), _) if partitionExpressionsColumns(opA.expressions) === Seq("a", "b") && partitionExpressionsColumns(opB.expressions) === Seq("a", "b") && opA.numPartitions == numPartitions && opB.numPartitions == numPartitions => j From 351e03506e268bdb5713d7a63caad3f3b87e62ad Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Wed, 22 Jul 2026 04:43:20 +0000 Subject: [PATCH 32/39] [SPARK-XXXXX][SS] Run stateful Real-Time Mode queries over a pipelined shuffle Enable a stateful streaming query (starting with deduplication) to run in Real-Time Mode by pipelining its repartition shuffle: - IncrementalExecution adds a Real-Time Mode preparation rule that marks every streaming shuffle exchange pipelined (ShuffleExchangeExec.pipelined = true). Real-Time Mode is detected structurally by a RealTimeStreamScanExec leaf. The producer (source scan) and consumer (stateful operator) stages are then co-scheduled and stream records through a transient shuffle instead of the consumer waiting for the producer to fully materialize. Inert for a non-RTM batch, so the ordinary microbatch path is unchanged. (AQE is already disabled for RTM batches upstream, so the pipelined shuffle is only ever submitted as a result job, never an AQE map-stage job.) - RealTimeModeAllowlist admits the operators a streaming dropDuplicates plans into (StreamingDeduplicateExec, StateStoreRestoreExec, StateStoreSaveExec) plus ShuffleExchangeExec, so the query is no longer rejected before it runs. Co-authored-by: Isaac --- .../runtime/IncrementalExecution.scala | 32 +++++++++++++++++-- .../runtime/RealTimeModeAllowlist.scala | 10 ++++++ 2 files changed, 40 insertions(+), 2 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala index 0d2e4a6941a00..0f2db3611208d 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala @@ -36,8 +36,9 @@ import org.apache.spark.sql.classic.SparkSession import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.execution.{CommandExecutionMode, LocalLimitExec, QueryExecution, SerializeFromObjectExec, SparkPlan, SparkPlanner, SparkStrategy => Strategy, UnaryExecNode} import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, MergingSessionsExec, ObjectHashAggregateExec, SortAggregateExec, UpdatingSessionsExec} +import org.apache.spark.sql.execution.datasources.v2.RealTimeStreamScanExec import org.apache.spark.sql.execution.datasources.v2.state.metadata.StateMetadataPartitionReader -import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike +import org.apache.spark.sql.execution.exchange.{ShuffleExchangeExec, ShuffleExchangeLike} import org.apache.spark.sql.execution.python.streaming.{FlatMapGroupsInPandasWithStateExec, TransformWithStateInPySparkExec} import org.apache.spark.sql.execution.streaming.{StreamingErrors, StreamingQueryPlanTraverseHelper} import org.apache.spark.sql.execution.streaming.checkpointing.{CheckpointFileManager, OffsetSeqMetadata, OffsetSeqMetadataBase} @@ -663,7 +664,34 @@ class IncrementalExecution( } } - override def preparations: Seq[Rule[SparkPlan]] = state +: super.preparations + /** + * For a Real-Time Mode batch, mark the shuffle exchanges as pipelined so the DAGScheduler + * co-schedules a stateful query's producer (source scan) and consumer (stateful operator) stages + * as one pipelined group -- records stream through a transient shuffle instead of the consumer + * waiting for the producer to fully materialize. The exchange carries the decision as a field + * (see ShuffleExchangeExec.pipelined); the PipelinedShuffleDependency it then builds is the whole + * opt-in -- routing to the streaming shuffle manager and pipelined-group scheduling both follow + * from that dependency type. + * + * Real-Time Mode is detected structurally by a RealTimeStreamScanExec leaf (there is no + * RTM-specific plan flag). Inert for a non-RTM batch, so the ordinary microbatch path is + * unchanged. + */ + object MarkPipelinedShuffleForRealTimeMode extends Rule[SparkPlan] { + override def apply(plan: SparkPlan): SparkPlan = { + val isRealTimeMode = plan.exists(_.isInstanceOf[RealTimeStreamScanExec]) + if (!isRealTimeMode) { + plan + } else { + plan.transformUp { + case s: ShuffleExchangeExec if !s.pipelined => s.copy(pipelined = true) + } + } + } + } + + override def preparations: Seq[Rule[SparkPlan]] = + state +: (super.preparations :+ MarkPipelinedShuffleForRealTimeMode) /** no need to try-catch again as this is already done once */ override def assertAnalyzed(): Unit = analyzed diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala index 7ae557797f36c..047ea3de4298d 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala @@ -55,7 +55,17 @@ object RealTimeModeAllowlist extends Logging { "org.apache.spark.sql.execution.datasources.v2.WriteToDataSourceV2Exec", "org.apache.spark.sql.execution.exchange.BroadcastExchangeExec", "org.apache.spark.sql.execution.exchange.ReusedExchangeExec", + // A pipelined (streaming) shuffle repartitions a stateful query's input by key. In Real-Time + // Mode the DAGScheduler co-schedules the shuffle's producer and consumer stages via a + // PipelinedShuffleDependency (see IncrementalExecution's pipelined-shuffle rule), so the + // exchange is a supported member of a pipelined group rather than a materialization barrier. + "org.apache.spark.sql.execution.exchange.ShuffleExchangeExec", "org.apache.spark.sql.execution.joins.BroadcastHashJoinExec", + // Streaming deduplication and the state-store access operators it plans into. These run in the + // pipelined-shuffle consumer stage, keyed by the same columns the shuffle repartitions on. + "org.apache.spark.sql.execution.streaming.operators.stateful.StateStoreRestoreExec", + "org.apache.spark.sql.execution.streaming.operators.stateful.StateStoreSaveExec", + "org.apache.spark.sql.execution.streaming.operators.stateful.StreamingDeduplicateExec", classOf[EventTimeWatermarkExec].getName ) From 4264a6ce87c4df6e07636cd92d79f052a4380a15 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Wed, 22 Jul 2026 04:50:25 +0000 Subject: [PATCH 33/39] [SPARK-XXXXX][SS][TESTS] Test stateful dedup running in Real-Time Mode over a pipelined shuffle Add an end-to-end test: a streaming dropDuplicates query in Real-Time Mode whose repartition shuffle is pipelined. It asserts correct dedup output across two batches, that every shuffle exchange in the executed plan is pipelined, and that the producer and consumer stages are co-scheduled (>= 2 stages running at once). The Real-Time Mode operator allowlist check is left enabled, so the test also covers admission of the dedup plan's operators. Co-authored-by: Isaac --- .../RealTimeModePipelinedShuffleSuite.scala | 98 +++++++++++++++++++ 1 file changed, 98 insertions(+) create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala new file mode 100644 index 0000000000000..4ddeda78ee037 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala @@ -0,0 +1,98 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.streaming + +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.atomic.AtomicInteger + +import org.apache.spark.scheduler.{SparkListener, SparkListenerStageCompleted} +import org.apache.spark.scheduler.SparkListenerStageSubmitted +import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec +import org.apache.spark.sql.execution.streaming.sources.{ContinuousMemorySink, LowLatencyMemoryStream} + +/** + * Tests that a stateful Real-Time Mode query (streaming `dropDuplicates`) runs when its repartition + * shuffle is a PipelinedShuffleDependency, so the source-scan producer stage and the dedup consumer + * stage are co-scheduled and stream records through a transient shuffle instead of the consumer + * waiting for the producer to fully materialize. + * + * The Real-Time Mode operator allowlist check is left enabled, so the test also confirms that the + * dedup plan's operators (StreamingDeduplicateExec, StateStoreRestore/Save, ShuffleExchangeExec) + * are admitted. It drives the source with the standard `testStream` DSL, which advances the query + * across multiple batches. + */ +class RealTimeModePipelinedShuffleSuite extends StreamRealTimeModeSuiteBase { + import testImplicits._ + + test("stateful dedup runs in Real-Time Mode over a pipelined shuffle") { + // Track, from the driver, whether the producer (source scan) and consumer (dedup) stages of the + // pipelined group were ever RUNNING simultaneously. runningStages holds currently-running stage + // ids (added on submit, removed on completion); maxConcurrentStages records the peak. A + // sequential producer-then-consumer schedule never exceeds one running stage at a time. + val runningStages = ConcurrentHashMap.newKeySet[Int]() + val maxConcurrentStages = new AtomicInteger(0) + val listener = new SparkListener { + override def onStageSubmitted(e: SparkListenerStageSubmitted): Unit = { + runningStages.add(e.stageInfo.stageId) + maxConcurrentStages.accumulateAndGet(runningStages.size(), Math.max) + } + override def onStageCompleted(e: SparkListenerStageCompleted): Unit = { + runningStages.remove(e.stageInfo.stageId) + } + } + spark.sparkContext.addSparkListener(listener) + + try { + val inputData = LowLatencyMemoryStream[(String, Int)] + + // scan --shuffle(repartition by key)--> streaming dropDuplicates --> sink. + // dropDuplicates on "key" forces a hash-partitioning ShuffleExchangeExec, which the + // IncrementalExecution rule marks pipelined for a Real-Time Mode batch. + val result = inputData + .toDF() + .select($"_1".as("key")) + .dropDuplicates("key") + .select($"key") + + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + // Batch 0: three distinct keys, each sent twice -> dedup emits each once. + AddData(inputData, ("a", 1), ("b", 1), ("c", 1), ("a", 2), ("b", 2), ("c", 2)), + StartStream(), + CheckAnswer("a", "b", "c"), + // Batch 1: all duplicates of already-seen keys plus one new key -> only "d" is new. + AddData(inputData, ("a", 3), ("b", 3), ("c", 3), ("d", 1)), + CheckAnswer("a", "b", "c", "d"), + // Every Real-Time Mode shuffle exchange in the executed plan is pipelined, and the producer + // + consumer stages were genuinely co-scheduled (>= 2 stages ran at once). + Execute { q => + val executedPlan = q.lastExecution.executedPlan + val exchanges = executedPlan.collect { case s: ShuffleExchangeExec => s } + assert(exchanges.nonEmpty, "expected a shuffle exchange in the dedup plan") + assert(exchanges.forall(_.pipelined), + "expected all Real-Time Mode shuffle exchanges to be pipelined, got: " + + exchanges.map(e => s"pipelined=${e.pipelined}").mkString(", ")) + assert(maxConcurrentStages.get() >= 2, + s"expected >= 2 stages running concurrently, saw max ${maxConcurrentStages.get()}") + }, + StopStream + ) + } finally { + spark.sparkContext.removeSparkListener(listener) + } + } +} From 2d503871fb967d67a761207cf9c1a951d28e8556 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Wed, 22 Jul 2026 05:08:32 +0000 Subject: [PATCH 34/39] [SPARK-XXXXX][SS][TESTS] Test dedup fault-tolerance in Real-Time Mode over a pipelined shuffle Add a fault-tolerance case to the pipelined-shuffle Real-Time Mode suite: a stateful dropDuplicates query whose repartition shuffle is pipelined, that hits a task failure mid-batch and then restarts from checkpoint. It asserts the query fails on the task error (Real-Time Mode does not retry tasks; a member failure aborts the whole co-scheduled group), and that on restart the deduplication state recovered from the last committed batch -- already-seen keys are not re-emitted. Switches the suite to the manual-clock base so batch boundaries are deterministic across the failure and restart, and uses CheckAnswerWithTimeout (which polls the sink in real time) with advanceRealTimeClock to drive Real-Time Mode batches. Co-authored-by: Isaac --- .../RealTimeModePipelinedShuffleSuite.scala | 78 +++++++++++++++++-- 1 file changed, 72 insertions(+), 6 deletions(-) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala index 4ddeda78ee037..e5fa82863d612 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala @@ -20,10 +20,17 @@ package org.apache.spark.sql.streaming import java.util.concurrent.ConcurrentHashMap import java.util.concurrent.atomic.AtomicInteger +import org.apache.spark.SparkException import org.apache.spark.scheduler.{SparkListener, SparkListenerStageCompleted} import org.apache.spark.scheduler.SparkListenerStageSubmitted import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec import org.apache.spark.sql.execution.streaming.sources.{ContinuousMemorySink, LowLatencyMemoryStream} +import org.apache.spark.sql.functions.udf + +/** Driver-side switch a UDF reads on executors to fail a task on demand (for fault-tolerance). */ +object RealTimeModePipelinedShuffleSuite { + @volatile var failTasks: Boolean = false +} /** * Tests that a stateful Real-Time Mode query (streaming `dropDuplicates`) runs when its repartition @@ -31,14 +38,19 @@ import org.apache.spark.sql.execution.streaming.sources.{ContinuousMemorySink, L * stage are co-scheduled and stream records through a transient shuffle instead of the consumer * waiting for the producer to fully materialize. * - * The Real-Time Mode operator allowlist check is left enabled, so the test also confirms that the + * The Real-Time Mode operator allowlist check is left enabled, so the tests also confirm that the * dedup plan's operators (StreamingDeduplicateExec, StateStoreRestore/Save, ShuffleExchangeExec) - * are admitted. It drives the source with the standard `testStream` DSL, which advances the query + * are admitted. They drive the source with the standard `testStream` DSL, which advances the query * across multiple batches. */ -class RealTimeModePipelinedShuffleSuite extends StreamRealTimeModeSuiteBase { +class RealTimeModePipelinedShuffleSuite extends StreamRealTimeModeManualClockSuiteBase { import testImplicits._ + override def beforeEach(): Unit = { + super.beforeEach() + RealTimeModePipelinedShuffleSuite.failTasks = false + } + test("stateful dedup runs in Real-Time Mode over a pipelined shuffle") { // Track, from the driver, whether the producer (source scan) and consumer (dedup) stages of the // pipelined group were ever RUNNING simultaneously. runningStages holds currently-running stage @@ -70,13 +82,17 @@ class RealTimeModePipelinedShuffleSuite extends StreamRealTimeModeSuiteBase { .select($"key") testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( - // Batch 0: three distinct keys, each sent twice -> dedup emits each once. + // Batch 0: three distinct keys, each sent twice -> dedup emits each once. A Real-Time Mode + // batch runs for a fixed duration and emits in real time, so CheckAnswerWithTimeout polls + // the sink (rather than blocking for batch completion) and advanceRealTimeClock ends it. AddData(inputData, ("a", 1), ("b", 1), ("c", 1), ("a", 2), ("b", 2), ("c", 2)), StartStream(), - CheckAnswer("a", "b", "c"), + CheckAnswerWithTimeout(60000, "a", "b", "c"), + advanceRealTimeClock, + WaitUntilBatchProcessed(0), // Batch 1: all duplicates of already-seen keys plus one new key -> only "d" is new. AddData(inputData, ("a", 3), ("b", 3), ("c", 3), ("d", 1)), - CheckAnswer("a", "b", "c", "d"), + CheckAnswerWithTimeout(60000, "a", "b", "c", "d"), // Every Real-Time Mode shuffle exchange in the executed plan is pipelined, and the producer // + consumer stages were genuinely co-scheduled (>= 2 stages ran at once). Execute { q => @@ -95,4 +111,54 @@ class RealTimeModePipelinedShuffleSuite extends StreamRealTimeModeSuiteBase { spark.sparkContext.removeSparkListener(listener) } } + + test("dedup over a pipelined shuffle recovers from a task failure via checkpoint restart") { + withTempDir { checkpointDir => + // A UDF that throws on demand, to fail a task mid-batch. Placed after the dedup so the + // failure lands in the pipelined consumer stage while the query is running. + val failUDF = udf { (key: String) => + if (RealTimeModePipelinedShuffleSuite.failTasks) { + throw new RuntimeException(s"forced task failure on $key") + } + key + } + + val inputData = LowLatencyMemoryStream[(String, Int)] + val result = inputData + .toDF() + .select($"_1".as("key")) + .dropDuplicates("key") + .select(failUDF($"key").as("key")) + + // First run: dedup "a","b" (batch 0 commits), then a task failure fails the query. RTM does + // not retry tasks, so a single task failure fails the whole batch/query. + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, ("a", 1), ("b", 1), ("a", 2)), + StartStream(checkpointLocation = checkpointDir.getAbsolutePath), + CheckAnswerWithTimeout(60000, "a", "b"), + advanceRealTimeClock, + WaitUntilBatchProcessed(0), + Execute { _ => RealTimeModePipelinedShuffleSuite.failTasks = true }, + AddData(inputData, ("c", 1)), + advanceRealTimeClock, + ExpectFailure[SparkException] { ex => + val msg = Option(ex.getCause).map(_.getMessage).getOrElse(ex.getMessage) + assert(msg != null && msg.contains("forced task failure"), + s"expected a forced task failure, got: $msg") + } + ) + + // Restart from the same checkpoint. The dedup state from batch 0 must survive: "a" and "b" + // are already seen and must NOT be re-emitted; only the genuinely-new "c","d" appear. This + // proves recovery-to-last-committed-batch works when the shuffle is pipelined. + RealTimeModePipelinedShuffleSuite.failTasks = false + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + StartStream(checkpointLocation = checkpointDir.getAbsolutePath), + AddData(inputData, ("a", 3), ("b", 3), ("c", 3), ("d", 1)), + CheckAnswerWithTimeout(60000, "c", "d"), + advanceRealTimeClock, + StopStream + ) + } + } } From e01c3af6c5abfb2637f27df62f7c060eda17cded Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Wed, 22 Jul 2026 06:33:13 +0000 Subject: [PATCH 35/39] [SPARK-XXXXX][SS][TESTS] Test dedup checkpoint-write-failure recovery in Real-Time Mode Add a second fault-tolerance case: a dedup query over a pipelined shuffle whose commit log write is fault-injected (via FailureInjectionCheckpointFileManager) so a batch cannot commit and the query fails. On restart from the same checkpoint, the committed batch's deduplication state survives and the uncommitted batch is re-run, so its data is still emitted. The injected close() failure surfaces as the underlying IOException in this path (the COMMIT_LOG_WRITE_FAILURE error-class wrapping is not open source), so the test expects that. Co-authored-by: Isaac --- .../RealTimeModePipelinedShuffleSuite.scala | 59 +++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala index e5fa82863d612..cf6abeff4884e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala @@ -17,6 +17,7 @@ package org.apache.spark.sql.streaming +import java.io.{File, IOException} import java.util.concurrent.ConcurrentHashMap import java.util.concurrent.atomic.AtomicInteger @@ -25,7 +26,9 @@ import org.apache.spark.scheduler.{SparkListener, SparkListenerStageCompleted} import org.apache.spark.scheduler.SparkListenerStageSubmitted import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec import org.apache.spark.sql.execution.streaming.sources.{ContinuousMemorySink, LowLatencyMemoryStream} +import org.apache.spark.sql.execution.streaming.state.{FailureInjectionCheckpointFileManager, FailureInjectionFileSystem, FailureInjectionState} import org.apache.spark.sql.functions.udf +import org.apache.spark.sql.internal.SQLConf /** Driver-side switch a UDF reads on executors to fail a task on demand (for fault-tolerance). */ object RealTimeModePipelinedShuffleSuite { @@ -161,4 +164,60 @@ class RealTimeModePipelinedShuffleSuite extends StreamRealTimeModeManualClockSui ) } } + + /** Run `f` with a temp dir whose checkpoint file ops can be fault-injected via injectionState. */ + private def withTempDirAllowFailureInjection(f: (File, FailureInjectionState) => Unit): Unit = { + withTempDir { dir => + val injectionState = FailureInjectionFileSystem.registerTempPath(dir.getPath) + try { + f(dir, injectionState) + } finally { + FailureInjectionFileSystem.removePathFromTempToInjectionState(dir.getPath) + } + } + } + + test("dedup over a pipelined shuffle recovers from a commit-log write failure") { + withSQLConf( + SQLConf.STREAMING_CHECKPOINT_FILE_MANAGER_CLASS.parent.key -> + classOf[FailureInjectionCheckpointFileManager].getName) { + withTempDirAllowFailureInjection { (checkpointDir, injectionState) => + val inputData = LowLatencyMemoryStream[(String, Int)] + val result = inputData + .toDF() + .select($"_1".as("key")) + .dropDuplicates("key") + .select($"key") + + // Batch 0 dedups "a","b" and commits. Then fail the close() of batch 1's commit-log write, + // so batch 1 cannot commit and the query fails after processing. + injectionState.createAtomicDelayCloseRegex = Seq(".*/commits/1") + + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, ("a", 1), ("b", 1)), + StartStream(checkpointLocation = checkpointDir.getAbsolutePath), + CheckAnswerWithTimeout(60000, "a", "b"), + advanceRealTimeClock, + WaitUntilBatchProcessed(0), + AddData(inputData, ("c", 1)), + CheckAnswerWithTimeout(60000, "a", "b", "c"), + advanceRealTimeClock, + // The injected close() failure surfaces as the underlying IOException failing the batch. + ExpectFailure[IOException]() + ) + + // Clear the injection and restart from the same checkpoint. Batch 0's dedup state must + // survive (a, b already seen); the uncommitted batch 1 is re-run, so its new key "c" is + // still emitted, and a further new key "d" is added. + injectionState.createAtomicDelayCloseRegex = Seq.empty + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + StartStream(checkpointLocation = checkpointDir.getAbsolutePath), + AddData(inputData, ("a", 2), ("b", 2), ("c", 2), ("d", 1)), + CheckAnswerWithTimeout(60000, "c", "d"), + advanceRealTimeClock, + StopStream + ) + } + } + } } From ab4234fcd81cacbeefcc68010f73aa0dc9a5e084 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 27 Jul 2026 19:15:46 +0000 Subject: [PATCH 36/39] [SPARK-XXXXX][CORE] Unregister a pipelined shuffle from the streaming output tracker on cleanup A pipelined shuffle is registered with the driver-only StreamingShuffleOutputTracker in DAGScheduler.createShuffleMapStage, alongside its MapOutputTracker registration, but nothing ever removed it. Over a long-running Real-Time Mode query -- which allocates a fresh shuffle id per micro-batch -- the tracker's shuffleInfos map grows without bound, and a re-registration of a reused id would hit the "registered twice" guard. Unregister it in ContextCleaner.doCleanupShuffle, right beside the existing mapOutputTrackerMaster.unregisterShuffle. The streaming tracker's state is driver-only (the worker tracker caches nothing and just RPCs the master), so a direct driver-side call is the correct counterpart -- unlike the MapOutputTracker, there is no per-executor cache to invalidate, so this does not belong in the RemoveShuffle RPC handler. The call is guarded by a new containsShuffle predicate so regular (non-pipelined) shuffles, which are never registered here, are skipped without emitting a spurious "not registered" warning on every shuffle cleanup. Also correct the StreamingShuffleManager.unregisterShuffle comment, which previously claimed the tracker was unregistered in BlockManagerStorageEndpoint's RemoveShuffle handler -- untrue in OSS (that path routes to this no-op manager method). Tests: a containsShuffle unit test, plus two ContextCleanerSuite integration tests (a pipelined shuffle is unregistered on cleanup; a regular shuffle is left untouched and quiet). Both integration tests were revert-checked -- each fails when its respective production change (the unregister call, or the containsShuffle guard) is removed. Co-authored-by: Isaac --- .../org/apache/spark/ContextCleaner.scala | 13 ++++++ .../spark/StreamingShuffleOutputTracker.scala | 4 ++ .../streaming/StreamingShuffleManager.scala | 4 +- .../apache/spark/ContextCleanerSuite.scala | 43 +++++++++++++++++++ .../StreamingShuffleOutputTrackerSuite.scala | 14 ++++++ 5 files changed, 76 insertions(+), 2 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/ContextCleaner.scala b/core/src/main/scala/org/apache/spark/ContextCleaner.scala index 0b3c22a22cb46..547112982284a 100644 --- a/core/src/main/scala/org/apache/spark/ContextCleaner.scala +++ b/core/src/main/scala/org/apache/spark/ContextCleaner.scala @@ -242,6 +242,17 @@ private[spark] class ContextCleaner( // to find blocks served by the shuffle service on deallocated executors shuffleDriverComponents.removeShuffle(shuffleId, blocking) mapOutputTrackerMaster.unregisterShuffle(shuffleId) + // A pipelined shuffle is also registered with the driver-only + // StreamingShuffleOutputTracker (see DAGScheduler.createShuffleMapStage). Its state lives + // solely on the driver -- the worker tracker caches nothing -- so unregister it directly + // here, alongside the MapOutputTracker cleanup, rather than through the RemoveShuffle RPC. + // Guarded by containsShuffle so regular (non-pipelined) shuffles are skipped without a + // spurious "not registered" warning. + streamingShuffleOutputTrackerMaster.foreach { tracker => + if (tracker.containsShuffle(shuffleId)) { + tracker.unregisterShuffle(shuffleId) + } + } listeners.asScala.foreach(_.shuffleCleaned(shuffleId)) logDebug("Cleaned shuffle " + shuffleId) } else { @@ -308,6 +319,8 @@ private[spark] class ContextCleaner( private def broadcastManager = sc.env.broadcastManager private def mapOutputTrackerMaster = sc.env.mapOutputTracker.asInstanceOf[MapOutputTrackerMaster] + private def streamingShuffleOutputTrackerMaster: Option[StreamingShuffleOutputTrackerMaster] = + sc.env.streamingShuffleOutputTracker.map(_.asInstanceOf[StreamingShuffleOutputTrackerMaster]) } private object ContextCleaner { diff --git a/core/src/main/scala/org/apache/spark/StreamingShuffleOutputTracker.scala b/core/src/main/scala/org/apache/spark/StreamingShuffleOutputTracker.scala index 090be7f691d4b..9d6fc4f1936c8 100644 --- a/core/src/main/scala/org/apache/spark/StreamingShuffleOutputTracker.scala +++ b/core/src/main/scala/org/apache/spark/StreamingShuffleOutputTracker.scala @@ -230,6 +230,10 @@ private[spark] class StreamingShuffleOutputTrackerMaster(conf: SparkConf) } } + // Whether the given shuffle is registered. ContextCleaner uses this on the driver to skip + // regular (non-pipelined) shuffles, which are never registered here, before unregistering. + def containsShuffle(shuffleId: Int): Boolean = shuffleInfos.containsKey(shuffleId) + // for testing purposes private[spark] def getShuffleInfo(shuffleId: Int): Option[StreamingShuffleInfo] = { Option(shuffleInfos.get(shuffleId)) diff --git a/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManager.scala b/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManager.scala index 03ca257353c59..b83bd33552b90 100644 --- a/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManager.scala +++ b/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManager.scala @@ -113,8 +113,8 @@ private[spark] class StreamingShuffleManager extends PipelinedShuffleManager wit override def unregisterShuffle(shuffleId: Int): Boolean = { // No manager-side state to release here: the driver's StreamingShuffleOutputTracker is - // unregistered in BlockManagerStorageEndpoint's RemoveShuffle handler, and per-task writer - // and reader resources are released via task completion listeners. + // unregistered in ContextCleaner.doCleanupShuffle (its state is driver-only), and per-task + // writer and reader resources are released via task completion listeners. true } diff --git a/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala b/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala index bb2d7d5c4d8e4..7b00e4440cce8 100644 --- a/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala +++ b/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala @@ -20,6 +20,7 @@ package org.apache.spark import scala.collection.mutable.HashSet import scala.util.Random +import org.apache.logging.log4j.Level import org.scalatest.BeforeAndAfter import org.scalatest.concurrent.Eventually._ import org.scalatest.concurrent.PatienceConfiguration @@ -127,6 +128,48 @@ class ContextCleanerSuite extends ContextCleanerSuiteBase { assert(rdd.collect().toList.equals(collected)) } + test("cleanup shuffle unregisters a pipelined shuffle from the streaming output tracker") { + // A pipelined shuffle is registered in BOTH the MapOutputTracker and the driver-only + // StreamingShuffleOutputTracker (see DAGScheduler.createShuffleMapStage). doCleanupShuffle + // must unregister it from the streaming tracker too, mirroring the MapOutputTracker cleanup; + // otherwise the tracker grows without bound across micro-batches in Real-Time Mode. + val streamingTracker = sc.env.streamingShuffleOutputTracker.get + .asInstanceOf[StreamingShuffleOutputTrackerMaster] + val mapOutputTracker = sc.env.mapOutputTracker.asInstanceOf[MapOutputTrackerMaster] + + // Register a shuffle id in both trackers exactly as the scheduler does for a pipelined shuffle. + val shuffleId = 1000 + mapOutputTracker.registerShuffle(shuffleId, numMaps = 2, numReduces = 2) + streamingTracker.registerShuffle(shuffleId, numMaps = 2, numReduces = 2, jobId = 0) + assert(streamingTracker.containsShuffle(shuffleId)) + + cleaner.doCleanupShuffle(shuffleId, blocking = true) + + // The streaming tracker entry is gone, alongside the MapOutputTracker cleanup. + assert(!streamingTracker.containsShuffle(shuffleId)) + assert(!mapOutputTracker.containsShuffle(shuffleId)) + } + + test("cleanup shuffle leaves a regular shuffle's streaming tracker untouched and quiet") { + // A regular (non-pipelined) shuffle is never registered with the streaming tracker, so + // doCleanupShuffle must skip it via the containsShuffle guard -- without emitting the + // "attempting to unregister a shuffle that hasn't been registered" warning for every shuffle. + val streamingTracker = sc.env.streamingShuffleOutputTracker.get + .asInstanceOf[StreamingShuffleOutputTrackerMaster] + + val (rdd, shuffleDeps) = newRDDWithShuffleDependencies() + rdd.collect() + shuffleDeps.foreach(s => assert(!streamingTracker.containsShuffle(s.shuffleId))) + + val logAppender = new LogAppender("unregister streaming shuffle") + logAppender.setThreshold(Level.WARN) + withLogAppender(logAppender, level = Some(Level.WARN)) { + shuffleDeps.foreach(s => cleaner.doCleanupShuffle(s.shuffleId, blocking = true)) + } + assert(!logAppender.loggingEvents.exists( + _.getMessage.getFormattedMessage.contains("hasn't been registered"))) + } + test("cleanup broadcast") { val broadcast = newBroadcast() val tester = new CleanerTester(sc, broadcastIds = Seq(broadcast.id)) diff --git a/core/src/test/scala/org/apache/spark/StreamingShuffleOutputTrackerSuite.scala b/core/src/test/scala/org/apache/spark/StreamingShuffleOutputTrackerSuite.scala index 3c8d9b2f14e21..a9351630ad8db 100644 --- a/core/src/test/scala/org/apache/spark/StreamingShuffleOutputTrackerSuite.scala +++ b/core/src/test/scala/org/apache/spark/StreamingShuffleOutputTrackerSuite.scala @@ -270,6 +270,20 @@ class StreamingShuffleOutputTrackerSuite assert(tracker.getAllShuffleWriterTaskLocations(0).isEmpty) } + test("StreamingShuffleOutputTrackerMaster - containsShuffle reflects register/unregister") { + val conf = new SparkConf(false) + val tracker = newTrackerMaster(conf) + + // A shuffle that was never registered is absent. + assert(!tracker.containsShuffle(0)) + + tracker.registerShuffle(shuffleId = 0, numMaps = 1, numReduces = 1, jobId = 1) + assert(tracker.containsShuffle(0)) + + tracker.unregisterShuffle(0) + assert(!tracker.containsShuffle(0)) + } + test("StreamingShuffleOutputTrackerMaster - register writer before shuffle fails") { val conf = new SparkConf(false) val tracker = newTrackerMaster(conf) From e27c581a6968dd9fe3b463abec6f6cdb96f395a2 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Tue, 28 Jul 2026 02:10:35 +0000 Subject: [PATCH 37/39] [SPARK-XXXXX][CORE] Fully separate pipelined shuffles from the MapOutputTracker A pipelined shuffle produces no durable, addressable map output: its producer/consumer task-location routing lives entirely in the StreamingShuffleOutputTracker, its availability is tracked on the stage (pipelinedCompletedPartitions), and its reader locates writers through the streaming tracker. Yet createShuffleMapStage still registered an empty ShuffleStatus for it in the MapOutputTracker, purely to keep containsShuffle / unregisterShuffle bookkeeping uniform. That coupled the streaming tracker's lifecycle to MapOutputTracker state: registration and cleanup of the streaming entry were both nested inside MapOutputTracker.containsShuffle guards, so a future divergence between the two trackers could double-register (throwing "registered twice") or leak the streaming entry. Split the two trackers cleanly by dependency type, with no overlap: - createShuffleMapStage: a PipelinedShuffleDependency registers ONLY with the StreamingShuffleOutputTracker (self-guarded on that tracker's own containsShuffle); a regular shuffle registers ONLY with the MapOutputTracker, exactly as before. No empty MapOutputTracker shell is created for a pipelined shuffle. - ContextCleaner.doCleanupShuffle: two independent branches, each keyed on its own tracker's membership -- MapOutputTracker cleanup for a regular shuffle, a direct driver-side StreamingShuffleOutputTracker unregister for a pipelined one (no RemoveShuffle RPC: a streaming shuffle has no block-manager-served files). Neither branch depends on the other tracker's state. This is safe because every MapOutputTracker path a pipelined shuffle could reach is either already type-gated out (output registration, availability, FetchFailed invalidation which routes to a group-atomic abort, push-merge, indeterminate/checksum rollback, map-stage-job stats) or tolerant of the entry's absence (executor/host-loss strips iterate shuffleStatuses and simply skip it). Consumer task locality never reads the tracker for a shuffle dependency (getPreferredLocsInternal recurses only through narrow deps). Tests (ContextCleanerSuite): a pipelined shuffle registered ONLY in the streaming tracker is still cleaned up by doCleanupShuffle even though it is absent from the MapOutputTracker (revert-checked: fails without the streaming cleanup branch); a regular shuffle takes the MapOutputTracker branch and never touches the streaming tracker. Co-authored-by: Isaac --- .../org/apache/spark/ContextCleaner.scala | 23 +++++----- .../apache/spark/scheduler/DAGScheduler.scala | 42 ++++++++++++------- .../spark/scheduler/ShuffleMapStage.scala | 13 +++--- .../apache/spark/ContextCleanerSuite.scala | 26 +++++++----- 4 files changed, 59 insertions(+), 45 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/ContextCleaner.scala b/core/src/main/scala/org/apache/spark/ContextCleaner.scala index 547112982284a..a01a480147d72 100644 --- a/core/src/main/scala/org/apache/spark/ContextCleaner.scala +++ b/core/src/main/scala/org/apache/spark/ContextCleaner.scala @@ -236,25 +236,26 @@ private[spark] class ContextCleaner( /** Perform shuffle cleanup. */ def doCleanupShuffle(shuffleId: Int, blocking: Boolean): Unit = { try { + // A shuffle lives in exactly one tracker, split by dependency type: a regular shuffle in the + // MapOutputTracker, a pipelined shuffle in the driver-only StreamingShuffleOutputTracker (see + // DAGScheduler.createShuffleMapStage). Clean up whichever holds it -- the two branches are + // independent, each keyed on its own tracker's membership, so neither depends on the other. if (mapOutputTrackerMaster.containsShuffle(shuffleId)) { logDebug("Cleaning shuffle " + shuffleId) // Shuffle must be removed before it's unregistered from the output tracker // to find blocks served by the shuffle service on deallocated executors shuffleDriverComponents.removeShuffle(shuffleId, blocking) mapOutputTrackerMaster.unregisterShuffle(shuffleId) - // A pipelined shuffle is also registered with the driver-only - // StreamingShuffleOutputTracker (see DAGScheduler.createShuffleMapStage). Its state lives - // solely on the driver -- the worker tracker caches nothing -- so unregister it directly - // here, alongside the MapOutputTracker cleanup, rather than through the RemoveShuffle RPC. - // Guarded by containsShuffle so regular (non-pipelined) shuffles are skipped without a - // spurious "not registered" warning. - streamingShuffleOutputTrackerMaster.foreach { tracker => - if (tracker.containsShuffle(shuffleId)) { - tracker.unregisterShuffle(shuffleId) - } - } listeners.asScala.foreach(_.shuffleCleaned(shuffleId)) logDebug("Cleaned shuffle " + shuffleId) + } else if (streamingShuffleOutputTrackerMaster.exists(_.containsShuffle(shuffleId))) { + // A pipelined shuffle's state is driver-only (the worker tracker caches nothing), so + // cleanup is a direct driver-side unregister. No RemoveShuffle RPC is needed: a streaming + // shuffle has no durable, block-manager-served files to remove. + logDebug("Cleaning pipelined shuffle " + shuffleId) + streamingShuffleOutputTrackerMaster.foreach(_.unregisterShuffle(shuffleId)) + listeners.asScala.foreach(_.shuffleCleaned(shuffleId)) + logDebug("Cleaned pipelined shuffle " + shuffleId) } else { logDebug("Asked to cleanup non-existent shuffle (maybe it was already removed)") } diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index d664a9fb3d96d..6dee4f5356f2e 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -703,7 +703,32 @@ private[spark] class DAGScheduler( shuffleIdToMapStage(shuffleDep.shuffleId) = stage updateJobIdStageIdMaps(jobId, stage) - if (!mapOutputTracker.containsShuffle(shuffleDep.shuffleId)) { + // A pipelined shuffle and a regular shuffle live in DIFFERENT trackers -- the split is by + // dependency type, with no overlap. A pipelined shuffle produces no durable, addressable map + // output; its producer/consumer task-location routing lives entirely in the + // StreamingShuffleOutputTracker, and it is never registered with the MapOutputTracker (nothing + // reads such an entry -- availability is tracked on the stage via pipelinedCompletedPartitions, + // the reader locates writers through the streaming tracker, and every MapOutputTracker path a + // pipelined shuffle could reach is either type-gated out or tolerant of its absence). A regular + // shuffle is registered with the MapOutputTracker exactly as before. + if (shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) { + // Register the producer/consumer fan-in before any consumer task asks for a writer location. + // Self-guarded on the streaming tracker's own state (createShuffleMapStage runs once per + // shuffleId via getOrCreateShuffleMapStage, but guard defensively against a re-entry). + sc.env.streamingShuffleOutputTracker.foreach { tracker => + val streamingTracker = tracker.asInstanceOf[StreamingShuffleOutputTrackerMaster] + if (!streamingTracker.containsShuffle(shuffleDep.shuffleId)) { + logInfo(log"Registering RDD ${MDC(RDD_ID, rdd.id)} " + + log"(${MDC(CREATION_SITE, rdd.getCreationSite)}) as input to pipelined " + + log"shuffle ${MDC(SHUFFLE_ID, shuffleDep.shuffleId)}") + streamingTracker.registerShuffle( + shuffleDep.shuffleId, + rdd.partitions.length, + shuffleDep.partitioner.numPartitions, + jobId) + } + } + } else if (!mapOutputTracker.containsShuffle(shuffleDep.shuffleId)) { // Kind of ugly: need to register RDDs with the cache and map output tracker here // since we can't do it in the RDD constructor because # of partitions is unknown logInfo(log"Registering RDD ${MDC(RDD_ID, rdd.id)} " + @@ -711,21 +736,6 @@ private[spark] class DAGScheduler( log"shuffle ${MDC(SHUFFLE_ID, shuffleDep.shuffleId)}") mapOutputTracker.registerShuffle(shuffleDep.shuffleId, rdd.partitions.length, shuffleDep.partitioner.numPartitions) - // A pipelined shuffle streams records from a live producer to a co-scheduled consumer, so the - // consumer must be able to look up its producer tasks' locations while both stages run. That - // routing lives in the StreamingShuffleOutputTracker (the MapOutputTracker registration above - // only tracks the empty, never-materialized durable output). Register here, inside the same - // once-per-shuffleId guard, so the tracker knows the producer/consumer fan-in before any - // consumer task asks for a writer location. Inert (tracker absent) for a regular shuffle. - if (shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) { - sc.env.streamingShuffleOutputTracker.foreach { tracker => - tracker.asInstanceOf[StreamingShuffleOutputTrackerMaster].registerShuffle( - shuffleDep.shuffleId, - rdd.partitions.length, - shuffleDep.partitioner.numPartitions, - jobId) - } - } } stage } diff --git a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala index 94078d698cf4f..135e5e5a4bf2c 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala @@ -64,13 +64,12 @@ private[spark] class ShuffleMapStage( /** * Availability tracking for a pipelined shuffle. A pipelined shuffle produces no durable, - * addressable map output. `createShuffleMapStage` still calls `MapOutputTracker.registerShuffle` - * for it (it is a `ShuffleDependency`, and an empty shuffle-status entry keeps the tracker's - * `containsShuffle` / `unregisterShuffle` bookkeeping uniform), but no map output is ever - * registered into that entry -- `registerMapOutput` runs only on the non-pipelined completion - * path -- so the entry stays empty. Its map-stage availability therefore cannot be read from the - * tracker; instead we track completed partitions here, monotonically: a partition is added when - * its map task succeeds and is NEVER removed on executor/host loss. + * addressable map output and is NOT registered with the `MapOutputTracker` at all -- + * `createShuffleMapStage` registers it only with the `StreamingShuffleOutputTracker` (the two + * trackers are split by dependency type, with no overlap). Its map-stage availability therefore + * cannot be read from the `MapOutputTracker`; instead we track completed partitions here, + * monotonically: a partition is added when its map task succeeds and is NEVER removed on + * executor/host loss. * * This is the crux of avoiding the streaming-writer resubmit hang: if a pipelined shuffle's * availability were read from the `MapOutputTracker`, losing an executor that held a completed diff --git a/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala b/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala index 7b00e4440cce8..e8f33b7803da9 100644 --- a/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala +++ b/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala @@ -129,31 +129,35 @@ class ContextCleanerSuite extends ContextCleanerSuiteBase { } test("cleanup shuffle unregisters a pipelined shuffle from the streaming output tracker") { - // A pipelined shuffle is registered in BOTH the MapOutputTracker and the driver-only - // StreamingShuffleOutputTracker (see DAGScheduler.createShuffleMapStage). doCleanupShuffle - // must unregister it from the streaming tracker too, mirroring the MapOutputTracker cleanup; - // otherwise the tracker grows without bound across micro-batches in Real-Time Mode. + // A pipelined shuffle lives ONLY in the driver-only StreamingShuffleOutputTracker -- it is + // never registered with the MapOutputTracker (see DAGScheduler.createShuffleMapStage). So the + // MapOutputTracker.containsShuffle guard is false for it, and doCleanupShuffle must still clean + // it up via its own streaming branch; otherwise the tracker grows without bound across + // micro-batches in Real-Time Mode. val streamingTracker = sc.env.streamingShuffleOutputTracker.get .asInstanceOf[StreamingShuffleOutputTrackerMaster] val mapOutputTracker = sc.env.mapOutputTracker.asInstanceOf[MapOutputTrackerMaster] - // Register a shuffle id in both trackers exactly as the scheduler does for a pipelined shuffle. + // Register a shuffle id in the streaming tracker ONLY, exactly as the scheduler now does for a + // pipelined shuffle (nothing is put in the MapOutputTracker). val shuffleId = 1000 - mapOutputTracker.registerShuffle(shuffleId, numMaps = 2, numReduces = 2) streamingTracker.registerShuffle(shuffleId, numMaps = 2, numReduces = 2, jobId = 0) assert(streamingTracker.containsShuffle(shuffleId)) + assert(!mapOutputTracker.containsShuffle(shuffleId), + "a pipelined shuffle must not be registered with the MapOutputTracker") cleaner.doCleanupShuffle(shuffleId, blocking = true) - // The streaming tracker entry is gone, alongside the MapOutputTracker cleanup. + // The streaming tracker entry is gone -- cleanup fired even though the shuffle was never in the + // MapOutputTracker, proving the streaming cleanup branch is independent of MapOutputTracker. assert(!streamingTracker.containsShuffle(shuffleId)) - assert(!mapOutputTracker.containsShuffle(shuffleId)) } test("cleanup shuffle leaves a regular shuffle's streaming tracker untouched and quiet") { - // A regular (non-pipelined) shuffle is never registered with the streaming tracker, so - // doCleanupShuffle must skip it via the containsShuffle guard -- without emitting the - // "attempting to unregister a shuffle that hasn't been registered" warning for every shuffle. + // A regular (non-pipelined) shuffle is registered only with the MapOutputTracker, never the + // streaming tracker. doCleanupShuffle must take the MapOutputTracker branch and never touch the + // streaming tracker -- in particular it must not emit the streaming tracker's "attempting to + // unregister a shuffle that hasn't been registered" warning. val streamingTracker = sc.env.streamingShuffleOutputTracker.get .asInstanceOf[StreamingShuffleOutputTrackerMaster] From b64c3933607a2388aacc6f4b2b47fa1c5816b54b Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Tue, 28 Jul 2026 02:48:45 +0000 Subject: [PATCH 38/39] [SPARK-XXXXX][CORE] Harden the MapOutputTracker decoupling: fail loud (before any state mutation) with no streaming tracker, and abort the group on a base-path pipelined FetchFailed Three hardening fixes to "Fully separate pipelined shuffles from the MapOutputTracker", all found by review of that change: 1. createShuffleMapStage registered a pipelined dependency only inside a streamingShuffleOutputTracker.foreach, with the MapOutputTracker registration in the else-if branch. If the streaming tracker were absent (a non-streaming custom incremental manager), the shuffle was registered in NEITHER tracker -- a consumer would then fail with an opaque NullPointerException in the shuffle writer. Fail loud instead: a PipelinedShuffleDependency requires a StreamingShuffleOutputTracker. Crucially, resolve the tracker UP FRONT -- before createShuffleMapStage mutates stageIdToStage / shuffleIdToMapStage / updateJobIdStageIdMaps -- so the fail-loud throw leaves no partial scheduler state. (handleJobSubmitted's catch does no rollback, so a throw after those mutations would leak a half-created stage; a re-submission would then reuse the stale cached stage and fail with a misleading cross-job-reuse error instead of re-throwing the real misconfiguration.) 2. A pipelined-group FetchFailed is normally intercepted by the group-atomic abort branch, whose guard is keyed on the failed stage's active job carrying hasPipelinedDependency. If a pipelined-shuffle FetchFailed reaches the base (non-group-atomic) path instead, that path is invalid for a pipelined shuffle in two ways: (a) its MapOutputTracker invalidation would throw ShuffleStatusNotFoundException (a pipelined shuffle is no longer registered there); and (b) its resubmit branch would enqueue a lone-stage resubmit of the transient producer, which cannot be re-read and would deadlock the group. Guard the whole base-path tail on mapStage.isPipelined and abort the group instead -- matching the group-atomic branch's outcome. Tests (revert-checked): - PipelinedShuffleRoutingSuite: a job with a pipelined dependency and no streaming tracker fails loud with the IllegalStateException, and a re-submission re-throws the SAME error -- proving no partial stage leaked (revert-checked against the throw-after-mutation placement, which instead leaks and yields a cross-job-reuse error on re-submit). - DAGSchedulerSuite: a FetchFailed for a pipelined shuffle reaching the base path (group-atomic guard forced to miss) aborts the group -- the job fails, the throwing MapOutputTracker calls are never invoked, and the transient producer is not resubmitted (revert-checked against both the no-guard and the invalidation-only-guard variants). Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 221 ++++++++++-------- .../spark/scheduler/DAGSchedulerSuite.scala | 53 +++++ .../PipelinedShuffleRoutingSuite.scala | 40 ++++ 3 files changed, 220 insertions(+), 94 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index 6dee4f5356f2e..7b07f62ade832 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -692,6 +692,24 @@ private[spark] class DAGScheduler( checkBarrierStageWithNumSlots(rdd, resourceProfile) checkBarrierStageWithRDDChainPattern(rdd, rdd.getNumPartitions) checkPipelinedProducerSupported(shuffleDep) + // A pipelined shuffle is served solely by the StreamingShuffleOutputTracker, so that tracker + // MUST exist for one -- it is created whenever a streaming-capable shuffle manager is set up + // (see SparkEnv.initializeStreamingShuffleOutputTracker), a prerequisite for producing a + // PipelinedShuffleDependency. Resolve it up front (BEFORE any stage-map mutation below), so a + // fail-loud on the misconfiguration leaves no partial scheduler state -- and so a pipelined + // shuffle is never silently registered in no tracker (a consumer would then find no writer + // locations). The reader enforces the same invariant (see StreamingShuffleReader). None for a + // regular shuffle, which is registered with the MapOutputTracker instead. + val streamingTrackerForPipelined: Option[StreamingShuffleOutputTrackerMaster] = + if (shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) { + Some(sc.env.streamingShuffleOutputTracker + .getOrElse(throw new IllegalStateException( + s"A pipelined shuffle (id ${shuffleDep.shuffleId}) requires a " + + "StreamingShuffleOutputTracker, but none is configured")) + .asInstanceOf[StreamingShuffleOutputTrackerMaster]) + } else { + None + } val numTasks = rdd.partitions.length val parents = getOrCreateParentStages(shuffleDeps, jobId) val id = nextStageId.getAndIncrement() @@ -706,17 +724,16 @@ private[spark] class DAGScheduler( // A pipelined shuffle and a regular shuffle live in DIFFERENT trackers -- the split is by // dependency type, with no overlap. A pipelined shuffle produces no durable, addressable map // output; its producer/consumer task-location routing lives entirely in the - // StreamingShuffleOutputTracker, and it is never registered with the MapOutputTracker (nothing - // reads such an entry -- availability is tracked on the stage via pipelinedCompletedPartitions, - // the reader locates writers through the streaming tracker, and every MapOutputTracker path a - // pipelined shuffle could reach is either type-gated out or tolerant of its absence). A regular - // shuffle is registered with the MapOutputTracker exactly as before. - if (shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) { - // Register the producer/consumer fan-in before any consumer task asks for a writer location. - // Self-guarded on the streaming tracker's own state (createShuffleMapStage runs once per - // shuffleId via getOrCreateShuffleMapStage, but guard defensively against a re-entry). - sc.env.streamingShuffleOutputTracker.foreach { tracker => - val streamingTracker = tracker.asInstanceOf[StreamingShuffleOutputTrackerMaster] + // StreamingShuffleOutputTracker (resolved fail-loud above), and it is never registered with the + // MapOutputTracker (nothing reads such an entry -- availability is tracked on the stage via + // pipelinedCompletedPartitions, the reader locates writers through the streaming tracker, and + // every MapOutputTracker path a pipelined shuffle could reach is either type-gated out or + // tolerant of its absence). A regular shuffle registers with the MapOutputTracker as before. + streamingTrackerForPipelined match { + case Some(streamingTracker) => + // Register the producer/consumer fan-in before any consumer task asks for a writer loc. + // Self-guarded on the tracker's own state (createShuffleMapStage runs once per shuffleId + // via getOrCreateShuffleMapStage, but guard defensively against a re-entry). if (!streamingTracker.containsShuffle(shuffleDep.shuffleId)) { logInfo(log"Registering RDD ${MDC(RDD_ID, rdd.id)} " + log"(${MDC(CREATION_SITE, rdd.getCreationSite)}) as input to pipelined " + @@ -727,15 +744,16 @@ private[spark] class DAGScheduler( shuffleDep.partitioner.numPartitions, jobId) } - } - } else if (!mapOutputTracker.containsShuffle(shuffleDep.shuffleId)) { - // Kind of ugly: need to register RDDs with the cache and map output tracker here - // since we can't do it in the RDD constructor because # of partitions is unknown - logInfo(log"Registering RDD ${MDC(RDD_ID, rdd.id)} " + - log"(${MDC(CREATION_SITE, rdd.getCreationSite)}) as input to " + - log"shuffle ${MDC(SHUFFLE_ID, shuffleDep.shuffleId)}") - mapOutputTracker.registerShuffle(shuffleDep.shuffleId, rdd.partitions.length, - shuffleDep.partitioner.numPartitions) + case None if !mapOutputTracker.containsShuffle(shuffleDep.shuffleId) => + // Kind of ugly: need to register RDDs with the cache and map output tracker here + // since we can't do it in the RDD constructor because # of partitions is unknown + logInfo(log"Registering RDD ${MDC(RDD_ID, rdd.id)} " + + log"(${MDC(CREATION_SITE, rdd.getCreationSite)}) as input to " + + log"shuffle ${MDC(SHUFFLE_ID, shuffleDep.shuffleId)}") + mapOutputTracker.registerShuffle(shuffleDep.shuffleId, rdd.partitions.length, + shuffleDep.partitioner.numPartitions) + + case _ => // regular shuffle already registered with the MapOutputTracker; nothing to do } stage } @@ -3461,86 +3479,101 @@ private[spark] class DAGScheduler( "longer running") } - if (mapStage.rdd.isBarrier()) { - // Mark all the map as broken in the map stage, to ensure retry all the tasks on - // resubmitted stage attempt. - // TODO: SPARK-35547: Clean all push-based shuffle metadata like merge enabled and - // TODO: finalized as we are clearing all the merge results. - mapOutputTracker.unregisterAllMapAndMergeOutput(shuffleId) - } else if (mapIndex != -1) { - // Mark the map whose fetch failed as broken in the map stage - mapOutputTracker.unregisterMapOutput(shuffleId, mapIndex, bmAddress) - if (pushBasedShuffleEnabled) { - // Possibly unregister the merge result , if the FetchFailed - // mapIndex is part of the merge result of - mapOutputTracker. - unregisterMergeResult(shuffleId, reduceId, bmAddress, Option(mapIndex)) - } + if (mapStage.isPipelined) { + // Defense-in-depth for a pipelined-shuffle FetchFailed that reaches this base path + // rather than the group-atomic branch above (e.g. a job whose hasPipelinedDependency + // flag was not propagated through the group check). The base path is invalid for a + // pipelined shuffle in two ways: (a) the MapOutputTracker invalidation below would + // throw ShuffleStatusNotFoundException (a pipelined shuffle is never registered there, + // see createShuffleMapStage); (b) the resubmit branch would enqueue a lone-stage + // resubmit of the transient producer, which -- as the group-atomic branch's comment + // explains -- is never valid and would deadlock the group. So abort the group here + // instead, matching the group-atomic branch's outcome, and skip both. + abortStage(failedStage, + s"A pipelined group member failed with a fetch failure: $failureMessage", None) } else { - // Unregister the merge result of if there is a FetchFailed event - // and is not a MetaDataFetchException which is signified by bmAddress being null - if (bmAddress != null && - bmAddress.executorId.equals(BlockManagerId.SHUFFLE_MERGER_IDENTIFIER)) { - assert(pushBasedShuffleEnabled, "Push based shuffle expected to " + - "be enabled when handling merge block fetch failure.") - mapOutputTracker. - unregisterMergeResult(shuffleId, reduceId, bmAddress, None) + if (mapStage.rdd.isBarrier()) { + // Mark all the map as broken in the map stage, to ensure retry all the tasks on + // resubmitted stage attempt. + // TODO: SPARK-35547: Clean all push-based shuffle metadata like merge enabled and + // TODO: finalized as we are clearing all the merge results. + mapOutputTracker.unregisterAllMapAndMergeOutput(shuffleId) + } else if (mapIndex != -1) { + // Mark the map whose fetch failed as broken in the map stage + mapOutputTracker.unregisterMapOutput(shuffleId, mapIndex, bmAddress) + if (pushBasedShuffleEnabled) { + // Possibly unregister the merge result , if the FetchFailed + // mapIndex is part of the merge result of + mapOutputTracker. + unregisterMergeResult(shuffleId, reduceId, bmAddress, Option(mapIndex)) + } + } else { + // Unregister the merge result of if there is a FetchFailed + // event and is not a MetaDataFetchException (signified by bmAddress being null) + if (bmAddress != null && + bmAddress.executorId.equals(BlockManagerId.SHUFFLE_MERGER_IDENTIFIER)) { + assert(pushBasedShuffleEnabled, "Push based shuffle expected to " + + "be enabled when handling merge block fetch failure.") + mapOutputTracker. + unregisterMergeResult(shuffleId, reduceId, bmAddress, None) + } } - } - - if (failedStage.rdd.isBarrier()) { - failedStage match { - case failedMapStage: ShuffleMapStage => - // Mark all the map as broken in the map stage, to ensure retry all the tasks on - // resubmitted stage attempt. - mapOutputTracker.unregisterAllMapAndMergeOutput(failedMapStage.shuffleDep.shuffleId) - case failedResultStage: ResultStage => - // Abort the failed result stage since we may have committed output for some - // partitions. - val reason = "Could not recover from a failed barrier ResultStage. Most recent " + - s"failure reason: $failureMessage" - abortStage(failedResultStage, reason, None) + if (failedStage.rdd.isBarrier()) { + failedStage match { + case failedMapStage: ShuffleMapStage => + // Mark all the map as broken in the map stage, to ensure retry all the tasks on + // resubmitted stage attempt. + mapOutputTracker.unregisterAllMapAndMergeOutput( + failedMapStage.shuffleDep.shuffleId) + + case failedResultStage: ResultStage => + // Abort the failed result stage since we may have committed output for some + // partitions. + val reason = "Could not recover from a failed barrier ResultStage. Most recent " + + s"failure reason: $failureMessage" + abortStage(failedResultStage, reason, None) + } } - } - if (shouldAbortStage) { - abortStage(failedStage, abortReason.get, None) - } else { // update failedStages and make sure a ResubmitFailedStages event is enqueued - // TODO: Cancel running tasks in the failed stage -- cf. SPARK-17064 - val noResubmitEnqueued = !failedStages.contains(failedStage) - failedStages += failedStage - failedStages += mapStage - if (noResubmitEnqueued) { - // For statically indeterminate stages, trigger rollback early (here and in - // submitMissingTasks) rather than deferring to task completion. This is more - // efficient because it clears shuffle outputs before the retry is submitted, - // ensuring findMissingPartitions() returns all partitions. - // - // For runtime detection (checksum mismatch), rollback is triggered at task - // completion when the mismatch is discovered. - // - // The `rollbackCurrentStage = true` parameter ensures the failed map stage is - // included in the cleanup: clearing its shuffle outputs, marking old task results - // to be ignored, and creating a new shuffle merge state for the upcoming retry. - if (mapStage.isStaticallyIndeterminate && - !mapStage.shuffleDep.checksumMismatchFullRetryEnabled) { - rollbackSucceedingStages(mapStage, rollbackCurrentStage = true) - } + if (shouldAbortStage) { + abortStage(failedStage, abortReason.get, None) + } else { // update failedStages and make sure a ResubmitFailedStages event is enqueued + // TODO: Cancel running tasks in the failed stage -- cf. SPARK-17064 + val noResubmitEnqueued = !failedStages.contains(failedStage) + failedStages += failedStage + failedStages += mapStage + if (noResubmitEnqueued) { + // For statically indeterminate stages, trigger rollback early (here and in + // submitMissingTasks) rather than deferring to task completion. This is more + // efficient because it clears shuffle outputs before the retry is submitted, + // ensuring findMissingPartitions() returns all partitions. + // + // For runtime detection (checksum mismatch), rollback is triggered at task + // completion when the mismatch is discovered. + // + // The `rollbackCurrentStage = true` parameter ensures the failed map stage is + // included in the cleanup: clearing its shuffle outputs, marking old task results + // to be ignored, and creating a new shuffle merge state for the upcoming retry. + if (mapStage.isStaticallyIndeterminate && + !mapStage.shuffleDep.checksumMismatchFullRetryEnabled) { + rollbackSucceedingStages(mapStage, rollbackCurrentStage = true) + } - // We expect one executor failure to trigger many FetchFailures in rapid succession, - // but all of those task failures can typically be handled by a single resubmission of - // the failed stage. We avoid flooding the scheduler's event queue with resubmit - // messages by checking whether a resubmit is already in the event queue for the - // failed stage. If there is already a resubmit enqueued for a different failed - // stage, that event would also be sufficient to handle the current failed stage, but - // producing a resubmit for each failed stage makes debugging and logging a little - // simpler while not producing an overwhelming number of scheduler events. - logInfo( - log"Resubmitting ${MDC(STAGE, mapStage)} " + - log"(${MDC(STAGE_NAME, mapStage.name)}) and ${MDC(FAILED_STAGE, failedStage)} " + - log"(${MDC(FAILED_STAGE_NAME, failedStage.name)}) due to fetch failure") - scheduleResubmit() + // We expect one executor failure to trigger many FetchFailures in rapid succession, + // but all of those task failures can typically be handled by a single resubmission + // of the failed stage. We avoid flooding the scheduler's event queue with resubmit + // messages by checking whether a resubmit is already in the event queue for the + // failed stage. If there is already a resubmit enqueued for a different failed + // stage, that event would also be sufficient to handle the current failed stage, + // but producing a resubmit for each failed stage makes debugging and logging a + // little simpler while not producing an overwhelming number of scheduler events. + logInfo( + log"Resubmitting ${MDC(STAGE, mapStage)} " + + log"(${MDC(STAGE_NAME, mapStage.name)}) and ${MDC(FAILED_STAGE, failedStage)} " + + log"(${MDC(FAILED_STAGE_NAME, failedStage.name)}) due to fetch failure") + scheduleResubmit() + } } } diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index 05ba990e2d638..2c3ad57f8fa5c 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7646,6 +7646,59 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } + test("pipelined shuffle: a FetchFailed reaching the base path aborts the group instead of " + + "throwing or resubmitting the transient producer") { + // A pipelined-group FetchFailed is normally intercepted by the group-atomic branch, whose guard + // is keyed on the failed stage's active job carrying hasPipelinedDependency. If that flag is + // ever not set/propagated for a job reading a pipelined shuffle, the FetchFailed falls through + // to the base path. That base path is invalid for a pipelined shuffle in two ways: (a) its + // MapOutputTracker invalidation would throw ShuffleStatusNotFoundException (a pipelined shuffle + // is never registered there); (b) its resubmit branch would enqueue a lone-stage resubmit of + // the transient producer, which cannot be re-read and would deadlock the group. The base path + // must instead abort the group. We reproduce the guard-miss by clearing the flag on the live + // active job, then driving the FetchFailed. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + val producerTs = taskSets.head + complete(producerTs, Seq( + (Success, makeMapStatus("hostA", 2)), + (Success, makeMapStatus("hostB", 2)))) + val consumerTs = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + val taskSetsBeforeFetchFailure = taskSets.size + + // Force the group-atomic guard to miss: clear hasPipelinedDependency on the live active job so + // the FetchFailed reaches the base path with a pipelined mapStage. + scheduler.activeJobs.foreach(_.hasPipelinedDependency = false) + assert(scheduler.shuffleIdToMapStage(pipelinedDep.shuffleId).isPipelined, + "the fetched shuffle must be the pipelined producer's") + + runEvent(makeCompletionEvent( + consumerTs.tasks(0), + FetchFailed(makeBlockManagerId("hostA"), pipelinedDep.shuffleId, 0L, 0, 0, "ignored"), + null)) + + // (a) The throwing MapOutputTracker invalidation was skipped for the pipelined id. + verify(mapOutputTracker, never()).unregisterAllMapAndMergeOutput(pipelinedDep.shuffleId) + verify(mapOutputTracker, never()).unregisterMapOutput( + org.mockito.ArgumentMatchers.eq(pipelinedDep.shuffleId), + org.mockito.ArgumentMatchers.anyInt(), + org.mockito.ArgumentMatchers.any()) + // (b) The group was aborted, NOT resubmitted: the job failed and no lone-stage resubmit of the + // transient producer was enqueued or launched. + scheduler.resubmitFailedStages() + assert(failure != null, "the group must be aborted (job failed) on the base-path FetchFailed") + assert(taskSets.size === taskSetsBeforeFetchFailure, + "a pipelined producer must NOT be resubmitted as a lone stage") + assert(!scheduler.runningStages.exists(_.isInstanceOf[ShuffleMapStage]), + "the transient pipelined producer must not be left running/resubmitted") + sc.listenerBus.waitUntilEmpty() + assertDataStructuresEmpty() + } + test("pipelined shuffle: a consumer result task throwing aborts the whole group") { // Defense-in-depth for the group-atomic model at the DAGScheduler layer: a CONSUMER (result) // task failing must tear down the whole group, not just its own stage. Result and map tasks diff --git a/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala b/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala index da74803a10ebc..f36649828ec86 100644 --- a/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala +++ b/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala @@ -240,6 +240,46 @@ class PipelinedShuffleRoutingSuite extends SparkFunSuite with LocalSparkContext assert(SparkEnv.get.streamingShuffleOutputTracker.isEmpty) } + test("createShuffleMapStage fails loud for a pipelined dependency when no streaming tracker") { + // A pipelined shuffle is served solely by the StreamingShuffleOutputTracker. With a non- + // streaming incremental manager configured, that tracker is absent (verified above), so a job + // whose stage graph contains a PipelinedShuffleDependency cannot be scheduled. + // createShuffleMapStage must fail loud rather than silently register the shuffle in no tracker + // (which would strand a consumer with no writer locations). This guards the decoupling's + // invariant that a pipelined dependency implies a streaming tracker. + // local[4]: the pipelined group's up-front slot admission (producer 2 + result 2 = 4) must pass + // so that stage creation is actually reached -- otherwise the job is rejected for slots first. + sc = new SparkContext("local[4]", "test", newConf()) + assert(SparkEnv.get.streamingShuffleOutputTracker.isEmpty) + val producer = sc.parallelize(1 to 4, 2).map(x => (x, x)) + val pipelined = new PipelinedShuffleDependency[Int, Int, Int](producer, new HashPartitioner(2)) + // A minimal reduce-side RDD whose single dependency is the pipelined shuffle, so submitting an + // action forces createShuffleMapStage for the pipelined producer. + val consumer = new RDD[(Int, Int)](sc, Seq(pipelined)) { + override def compute(split: Partition, ctx: TaskContext): Iterator[(Int, Int)] = + Iterator.empty + override protected def getPartitions: Array[Partition] = + Array.tabulate(2)(i => new Partition { override def index: Int = i }) + } + val ex = intercept[Exception] { + consumer.count() + } + assert(findCause[IllegalStateException](ex).exists( + _.getMessage.contains("requires a StreamingShuffleOutputTracker")), + s"expected a fail-loud IllegalStateException about the missing tracker, got: $ex") + + // The fail-loud throw happens BEFORE createShuffleMapStage mutates stageIdToStage / + // shuffleIdToMapStage, so it leaves no partial scheduler state behind. If it left a half-created + // stage cached in shuffleIdToMapStage, a re-submission would reuse that stale stage instead of + // re-throwing. Re-submitting the same job must therefore throw the SAME fail-loud error again. + val ex2 = intercept[Exception] { + consumer.count() + } + assert(findCause[IllegalStateException](ex2).exists( + _.getMessage.contains("requires a StreamingShuffleOutputTracker")), + s"a re-submission must re-throw the fail-loud error (no leaked partial stage), got: $ex2") + } + test("spark.shuffle.manager.incremental resolves the same short aliases as the default manager") { // "sort" is a short alias the incremental slot must resolve (to SortShuffleManager) rather than // treat as a class name. SortShuffleManager is blocking, so it is rejected from the pipelined From ae85d3f4dcb5b58fe860b99cc20d4ad8fe5dc02c Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Tue, 28 Jul 2026 23:45:45 +0000 Subject: [PATCH 39/39] [SPARK-XXXXX][SS] Fix scalastyle line length and document the RTM single-shuffle assumption Two small follow-ups from a review pass: - PipelinedShuffleRoutingSuite: reflow a 101-char comment line (dev/scalastyle flagged it as exceeding 100 characters). Comment-only, no behavior change. - IncrementalExecution.MarkPipelinedShuffleForRealTimeMode: document why marking every ShuffleExchangeExec pipelined is safe today (an RTM plan has exactly one shuffle -- a second is rejected by RealTimeModeAllowlist, see the "repartition not allowed" test / SPARK-54237), and record the latent hazard for when RTM is widened to multi-shuffle plans: transformUp cannot reach a ReusedExchangeExec's wrapped child (a leaf field, not a tree child), so a reused shuffle would stay regular while its twin becomes pipelined -- a mixed regular+pipelined job the scheduler rejects. Note the fix options (mark inside ReusedExchangeExec, run before exchange reuse, or disable reuse for RTM) for that future change. Co-authored-by: Isaac --- .../spark/shuffle/PipelinedShuffleRoutingSuite.scala | 6 +++--- .../streaming/runtime/IncrementalExecution.scala | 10 ++++++++++ 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala b/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala index f36649828ec86..8e0f9e04c1bd8 100644 --- a/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala +++ b/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala @@ -269,9 +269,9 @@ class PipelinedShuffleRoutingSuite extends SparkFunSuite with LocalSparkContext s"expected a fail-loud IllegalStateException about the missing tracker, got: $ex") // The fail-loud throw happens BEFORE createShuffleMapStage mutates stageIdToStage / - // shuffleIdToMapStage, so it leaves no partial scheduler state behind. If it left a half-created - // stage cached in shuffleIdToMapStage, a re-submission would reuse that stale stage instead of - // re-throwing. Re-submitting the same job must therefore throw the SAME fail-loud error again. + // shuffleIdToMapStage, so it leaves no partial scheduler state behind. If it left a + // half-created stage cached in shuffleIdToMapStage, a re-submission would reuse that stale + // stage instead of re-throwing. Re-submitting the same job must therefore re-throw the error. val ex2 = intercept[Exception] { consumer.count() } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala index 0f2db3611208d..a037369e3addb 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala @@ -676,6 +676,16 @@ class IncrementalExecution( * Real-Time Mode is detected structurally by a RealTimeStreamScanExec leaf (there is no * RTM-specific plan flag). Inert for a non-RTM batch, so the ordinary microbatch path is * unchanged. + * + * Single-shuffle assumption: today an RTM plan has exactly one shuffle exchange (a second one is + * rejected by RealTimeModeAllowlist -- see the "repartition not allowed" test / SPARK-54237), so + * transformUp marks that one exchange and there is no exchange reuse to worry about. If RTM is + * ever widened to multi-shuffle plans (stateful joins, CTEs read twice, a union of the same keyed + * stream), transformUp would NOT reach a ReusedExchangeExec's wrapped `child` (it is a leaf whose + * child is a field, not a tree child), so a reused shuffle would stay regular while its in-tree + * twin becomes pipelined -- a mixed regular+pipelined job the scheduler rejects every batch. When + * that day comes, mark inside ReusedExchangeExec (or run this before exchange reuse, or disable + * reuse for RTM) and add a multi-shuffle test. */ object MarkPipelinedShuffleForRealTimeMode extends Rule[SparkPlan] { override def apply(plan: SparkPlan): SparkPlan = {