From 6b780dff4e163ab118045c33ffd703381a3faa1e Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Fri, 10 Jul 2026 22:55:55 +0000 Subject: [PATCH 01/54] [SPARK-XXXXX][CORE] Concurrently schedule stages connected by a pipelined shuffle Native DAGScheduler support for concurrent-stage scheduling over a PipelinedShuffleDependency (from the prior PR). A pipelined shuffle is incrementally readable: its consumer stage may begin reading output while the producer is still running, so the two are co-scheduled ('pipelined group') instead of the consumer waiting for the producer to materialize. submitStage: when a stage has missing parents, classify them by their shuffle dependency type. A parent read through a PipelinedShuffleDependency is a pipelined parent. The stage is co-scheduled with its producers (its tasks submitted immediately) only if every missing parent is pipelined AND each is already running; otherwise it parks in waitingStages exactly as before. A stage with a regular missing parent, or a pipelined parent not yet running, waits and is reconsidered later. This is inert for jobs with no pipelined dependency -- the full DAGSchedulerSuite is unchanged. submitWaitingPipelinedChildStages: the 'producer started running' analog of submitWaitingChildStages. When a pipelined producer starts, its waiting consumers are co-scheduled immediately (not only when the producer completes), so a consumer parked because its producer sat behind a regular shuffle is co-scheduled as soon as the producer runs. Cascades transitively. handleJobSubmitted: reject a job that uses a pipelined dependency when speculation is enabled -- a speculative producer copy would race a consumer already reading the producer's partial output, with no commit barrier. The check runs before stage creation so a rejected job leaves no scheduler state behind. Tests (DAGSchedulerSuite): concurrent submission, inertness for a regular shuffle, mixed pipelined+regular parents, a deep all-pipelined chain, a pipelined producer behind a regular shuffle (must not co-schedule early), two pipelined parents (co-schedule and re-park semantics), fan-out reconsideration to multiple consumers, transitive cascade, no double-submission on producer completion, and speculation rejection (plus that a regular job under speculation is not rejected). Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 137 +++++- .../spark/scheduler/DAGSchedulerSuite.scala | 397 ++++++++++++++++++ 2 files changed, 533 insertions(+), 1 deletion(-) 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 5e608308de297..293cd73b4e383 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -945,6 +945,52 @@ private[spark] class DAGScheduler( missing.toList } + /** + * Whether the RDD graph rooted at `finalRDD` contains a [[PipelinedShuffleDependency]] anywhere. + * Walks the RDD dependency graph directly (not the stage graph), so it can be checked before any + * stages are created -- letting a job be rejected up front without leaving partial scheduler + * state behind. + */ + private def rddGraphHasPipelinedDependency(finalRDD: RDD[_]): Boolean = { + // traverseRDDGraphUntil stops and returns false as soon as the visitor returns false; we use + // that to short-circuit on the first pipelined dependency found. It returns true if the whole + // graph was visited without stopping (i.e. none found), so negate the result. + !traverseRDDGraphUntil(finalRDD) { (rdd, enqueue) => + val hasPipelined = rdd.dependencies.exists { + case _: PipelinedShuffleDependency[_, _, _] => true + case dep => + enqueue(dep.rdd) + false + } + !hasPipelined // keep traversing while none found; stop (return false) when one is found + } + } + + /** + * Reject a job that enables speculation and uses a pipelined shuffle: a speculative copy of a + * producer would race a consumer already reading the producer's partial output, with no commit + * barrier protecting the read (spec S9). Checked up front, before any stage is created, so a + * rejection leaves no partial scheduler state. Used by the result-job path (handleJobSubmitted); + * the map-stage-job path rejects a pipelined dependency outright (see handleMapStageSubmitted), + * which subsumes speculation. Returns true (and fails the job via `listener`) if rejected; false + * otherwise. Inert for jobs without a pipelined dependency. + */ + private def rejectSpeculationWithPipelinedShuffle( + jobId: Int, + finalRDD: RDD[_], + listener: JobListener): Boolean = { + if (sc.conf.get(config.SPECULATION_ENABLED) && rddGraphHasPipelinedDependency(finalRDD)) { + logWarning(log"Rejecting job ${MDC(JOB_ID, jobId)}: speculation is incompatible with " + + log"pipelined shuffle dependencies") + listener.jobFailed(new SparkException( + "Speculative execution is not supported for a job that uses a pipelined shuffle " + + s"dependency. Disable ${config.SPECULATION_ENABLED.key} for such jobs.")) + true + } else { + false + } + } + /** Invoke `.partitions` on the given RDD and all of its ancestors */ private def eagerlyComputePartitionsForRddAndAncestors(rdd: RDD[_]): Unit = { val startTime = System.nanoTime @@ -1346,6 +1392,40 @@ private[spark] class DAGScheduler( } } + /** + * Reconsider waiting stages that read `runningParent` through a pipelined edge, now that it has + * started running. A pipelined edge is non-sequencing, so such a consumer can be co-scheduled as + * soon as its producer is running -- it need not wait for the producer to finish. This is the + * "producer started" analog of `submitWaitingChildStages` (which fires on producer *completion*): + * without it, a pipelined consumer parked because its producer was not yet runnable (e.g. the + * producer sat behind a regular shuffle) would not be co-scheduled until the producer finished, + * losing the pipelining. Inert unless `runningParent` is the pipelined producer of a waiting + * stage. + */ + private def submitWaitingPipelinedChildStages(runningParent: Stage): Unit = { + val pipelinedChildren = waitingStages.filter { child => + child.parents.contains(runningParent) && isPipelinedProducer(runningParent) + }.toArray + // Remove them from waitingStages before resubmitting, or submitStage's `!waitingStages(stage)` + // guard would treat them as already-scheduled and no-op (mirrors submitWaitingChildStages). + waitingStages --= pipelinedChildren + for (child <- pipelinedChildren.sortBy(_.firstJobId)) { + logInfo(log"Reconsidering ${MDC(STAGE, child)} now that its pipelined producer " + + log"${MDC(STAGE, runningParent)} is running") + submitStage(child) + } + } + + /** + * Whether `stage` produces its output through a [[PipelinedShuffleDependency]] -- i.e. it is a + * pipelined producer, whose consumers may run concurrently with it. Callers that need the + * producer/consumer relationship check `child.parents.contains(stage)` separately. + */ + private def isPipelinedProducer(stage: Stage): Boolean = stage match { + case m: ShuffleMapStage => m.shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]] + case _ => false + } + /** Finds the earliest-created active job that needs the stage */ // TODO: Probably should actually find among the active jobs that need this // stage the one with the highest priority (highest-priority pool, earliest created). @@ -1489,6 +1569,15 @@ private[spark] class DAGScheduler( return } + // A job that uses a pipelined shuffle runs its producer and consumer stages concurrently, so a + // speculative copy of a producer would race a consumer already reading the producer's partial + // output, with no commit barrier protecting the read. Reject such a job up front -- before any + // stages are created, so no partial scheduler state is left behind. (Inert for jobs without + // any pipelined dependency.) + if (rejectSpeculationWithPipelinedShuffle(jobId, finalRDD, listener)) { + return + } + var finalStage: ResultStage = null try { // New stage creation may throw an exception if, for example, jobs are run on a @@ -1564,6 +1653,23 @@ private[spark] class DAGScheduler( listener: JobListener, artifacts: JobArtifactSet, properties: Properties): Unit = { + // A map-stage job (SparkContext.submitMapStage, e.g. AQE's ShuffleExchangeExec) materializes a + // shuffle to produce its map-output statistics. A PipelinedShuffleDependency cannot serve that: + // it is transient with no durable, addressable map output, so it registers no map outputs and + // getStatistics would return all-zero stats (misleading the map-stage/AQE consumer); and its + // producer/consumer co-scheduling makes speculation unsafe (spec S9). Reject a pipelined + // dependency submitted as a map-stage job outright, up front, before any stage is created so no + // partial scheduler state is left behind. Inert for a regular ShuffleDependency. (The + // result-job path, handleJobSubmitted, only rejects the speculation case, since a pipelined + // dependency is a legitimate internal edge there.) + if (dependency.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) { + logWarning(log"Rejecting map-stage job ${MDC(JOB_ID, jobId)}: a pipelined shuffle dependency " + + log"cannot be materialized as a map-stage job") + listener.jobFailed(new SparkException( + "A pipelined shuffle dependency cannot be submitted as a map-stage job: it has no durable " + + "map output to produce statistics from. This is not supported.")) + return + } // Submitting this map stage might still require the creation of some parent stages, so make // sure that happens. var finalStage: ShuffleMapStage = null @@ -1630,11 +1736,40 @@ private[spark] class DAGScheduler( logInfo(log"Submitting ${MDC(STAGE, stage)} (${MDC(RDD_ID, stage.rdd)}), " + log"which has no missing parents") submitMissingTasks(stage, jobId.get) + // If this stage is the pipelined producer of a waiting consumer, co-schedule it now. + submitWaitingPipelinedChildStages(stage) } else { for (parent <- missing) { submitStage(parent) } - waitingStages += stage + + // A missing parent reached through a PipelinedShuffleDependency ("pipelined parent") + // is incrementally readable: this stage may run before that parent materializes, so + // the two are co-scheduled. `missing` already holds the direct parent shuffle-map + // stages (from getMissingParentStages), so classify them by their shuffle dependency + // type -- no extra graph walk. For a job with no pipelined dependency, pipelinedMissing + // is empty and this stage simply parks in waitingStages, exactly as before. + val (pipelinedMissing, regularMissing) = missing.partition(isPipelinedProducer) + // Co-schedule only if EVERY missing parent is pipelined AND each is actually running + // now. submitStage above may have parked a pipelined parent in waitingStages (e.g. it + // has its own regular missing parent); running this stage against a not-yet-running + // producer would strand it. If any parent is regular or not yet runnable, park this + // stage. It is then resubmitted and co-scheduled once its parents become runnable -- + // via submitWaitingChildStages when a regular parent completes, or via + // submitWaitingPipelinedChildStages when a pipelined parent starts running. + val allPipelinedParentsRunning = + pipelinedMissing.nonEmpty && pipelinedMissing.forall(runningStages.contains) + + if (regularMissing.isEmpty && allPipelinedParentsRunning) { + logInfo(log"Submitting ${MDC(STAGE, stage)} concurrently with its running pipelined " + + log"producer(s) ${MDC(MISSING_PARENT_STAGES, pipelinedMissing)}") + submitMissingTasks(stage, jobId.get) + // This stage is now running; if it is itself the pipelined producer of a waiting + // consumer, co-schedule that consumer too. + submitWaitingPipelinedChildStages(stage) + } else { + waitingStages += stage + } } } } 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 e12348e1be2d7..a58e60de0325c 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6148,6 +6148,403 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } CompletionEvent(task, reason, result, accumUpdates ++ extraAccumUpdates, metricPeaks, taskInfo) } + + // ========================================================================================== + // Pipelined shuffle dependency: group formation + concurrent submission (M1.2) + // ========================================================================================== + + test("pipelined shuffle: consumer stage is submitted concurrently with its producer") { + // producer (shuffle map) --[pipelined]--> consumer (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) + submit(consumerRdd, Array(0, 1)) + + // Both the producer's and the consumer's task sets must be submitted before the producer + // completes: the pipelined edge is non-sequencing, so the consumer does not wait. + assert(taskSets.size === 2, + s"expected producer and consumer submitted concurrently, got ${taskSets.size} task sets") + val producerStageId = taskSets.head.stageId + val consumerStageId = taskSets(1).stageId + assert(producerStageId != consumerStageId) + // The consumer is actually running (not parked in waitingStages), co-resident with the + // producer -- both are in runningStages. + assert(!scheduler.waitingStages.exists(_.rdd eq consumerRdd), + "consumer should not be waiting; it is co-scheduled with its pipelined producer") + assert(scheduler.runningStages.exists(_.rdd eq consumerRdd)) + assert(scheduler.runningStages.exists(_.rdd eq producerRdd)) + + // Now complete the producer, then the consumer, and the job finishes. + completeShuffleMapStageSuccessfully(producerStageId, 0, 2) + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("regular shuffle still waits (pipelined change is inert without a pipelined dependency)") { + // Identical shape but a REGULAR shuffle dependency: the consumer must NOT be submitted until + // the producer materializes -- verifying the concurrent-submission path is inert here. + 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)) + + // Only the producer is submitted initially; the consumer waits. + assert(taskSets.size === 1, s"expected only the producer submitted, got ${taskSets.size}") + val producerStageId = taskSets.head.stageId + completeShuffleMapStageSuccessfully(producerStageId, 0, 2) + // Now the consumer is submitted. + assert(taskSets.size === 2) + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: mixed parents wait for the regular parent, co-schedule the pipelined") { + // consumer depends on a pipelined producer AND a regular producer. It must wait for the + // regular one, but when resubmitted it co-schedules with the (still-running) pipelined one. + val pipelinedProducerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(pipelinedProducerRdd, new HashPartitioner(2)) + val regularProducerRdd = new MyRDD(sc, 2, Nil) + val regularDep = new ShuffleDependency(regularProducerRdd, new HashPartitioner(2)) + val consumerRdd = + new MyRDD(sc, 2, List(pipelinedDep, regularDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + + // Both producers are submitted, but the consumer waits (it has a regular missing parent). + assert(taskSets.size === 2, s"expected both producers submitted, got ${taskSets.size}") + assert(scheduler.waitingStages.exists(_.rdd eq consumerRdd), + "consumer should be waiting on its regular parent") + + // Complete the regular producer. Now the consumer is resubmitted and, since its only remaining + // missing parent is pipelined, co-scheduled with it (a 3rd task set appears). + val regularStageId = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq regularProducerRdd + }.get.stageId + completeShuffleMapStageSuccessfully(regularStageId, 0, 2) + assert(taskSets.size === 3, "consumer should be co-scheduled once the regular parent finishes") + + // Finish the pipelined producer and the consumer. + val pipelinedStageId = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq pipelinedProducerRdd + }.get.stageId + completeShuffleMapStageSuccessfully(pipelinedStageId, 0, 2) + 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: 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) + val psdAB = new PipelinedShuffleDependency(rddA, new HashPartitioner(2)) + val rddB = new MyRDD(sc, 2, List(psdAB), tracker = mapOutputTracker) + val psdBC = new PipelinedShuffleDependency(rddB, new HashPartitioner(2)) + val rddC = new MyRDD(sc, 2, List(psdBC), tracker = mapOutputTracker) + submit(rddC, Array(0, 1)) + + // All three stages must be running concurrently; none parked. + assert(taskSets.size === 3, s"expected A, B, C co-scheduled, got ${taskSets.size}") + assert(scheduler.waitingStages.isEmpty, "no stage should be waiting in an all-pipelined chain") + Seq(rddA, rddB, rddC).foreach { rdd => + assert(scheduler.runningStages.exists(_.rdd eq rdd), s"stage for $rdd should be running") + } + + // Drain in dependency order: A, then B, then the result stage C. + val idA = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddA).get.stageId + completeShuffleMapStageSuccessfully(idA, 0, 2) + val idB = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddB).get.stageId + completeShuffleMapStageSuccessfully(idB, 0, 2) + val tsC = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddC).get + complete(tsC, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a pipelined producer behind a regular shuffle is NOT co-scheduled early") { + // A --regular--> B --pipelined--> C. B cannot run until A materializes, so C must NOT be + // co-scheduled with B while B is still parked. (Regression: the naive check co-scheduled C + // against a not-yet-running B, stranding C -- a hang under slot pressure.) + val rddA = new MyRDD(sc, 2, Nil) + val regAB = new ShuffleDependency(rddA, new HashPartitioner(2)) + val rddB = new MyRDD(sc, 2, List(regAB), tracker = mapOutputTracker) + val psdBC = new PipelinedShuffleDependency(rddB, new HashPartitioner(2)) + val rddC = new MyRDD(sc, 2, List(psdBC), tracker = mapOutputTracker) + submit(rddC, Array(0, 1)) + + // Only A runs initially. B waits on A; C waits on B (its pipelined parent isn't running yet). + assert(taskSets.size === 1, s"expected only A submitted, got ${taskSets.size}") + assert(scheduler.waitingStages.exists(_.rdd eq rddB), "B should wait on its regular parent A") + assert(scheduler.waitingStages.exists(_.rdd eq rddC), + "C must NOT be co-scheduled while its pipelined producer B is still parked") + + // A finishes -> B becomes runnable and is now co-scheduled with C (its pipelined consumer). + val idA = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddA).get.stageId + completeShuffleMapStageSuccessfully(idA, 0, 2) + assert(scheduler.runningStages.exists(_.rdd eq rddB), "B should now be running") + assert(scheduler.runningStages.exists(_.rdd eq rddC), + "C should now be co-scheduled with the running B") + + val idB = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddB).get.stageId + completeShuffleMapStageSuccessfully(idB, 0, 2) + val tsC = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddC).get + complete(tsC, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: consumer with two pipelined parents co-schedules with both") { + val rddA = new MyRDD(sc, 2, Nil) + val psdA = new PipelinedShuffleDependency(rddA, new HashPartitioner(2)) + val rddB = new MyRDD(sc, 2, Nil) + val psdB = new PipelinedShuffleDependency(rddB, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(psdA, psdB), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + + // Both producers and the consumer are all co-scheduled. + assert(taskSets.size === 3, s"expected both producers + consumer, got ${taskSets.size}") + assert(!scheduler.waitingStages.exists(_.rdd eq consumerRdd)) + Seq(rddA, rddB, consumerRdd).foreach { rdd => + assert(scheduler.runningStages.exists(_.rdd eq rdd)) + } + + val idA = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddA).get.stageId + completeShuffleMapStageSuccessfully(idA, 0, 2) + val idB = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddB).get.stageId + completeShuffleMapStageSuccessfully(idB, 0, 2) + val tsC = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd).get + complete(tsC, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: consumer with two pipelined parents re-parks until BOTH are running") { + // R1 --regular--> P1 --pipelined--> C ; R2 --regular--> P2 --pipelined--> C. + // When P1 starts (R1 done) but P2 is still parked (R2 not done), C must RE-PARK, not + // co-schedule -- else it would run against a not-yet-running P2. (Guards the `forall` in the + // co-schedule gate against a `.exists` regression.) + val rddR1 = new MyRDD(sc, 2, Nil) + val rddP1 = new MyRDD(sc, 2, List(new ShuffleDependency(rddR1, new HashPartitioner(2))), + tracker = mapOutputTracker) + val psdP1 = new PipelinedShuffleDependency(rddP1, new HashPartitioner(2)) + val rddR2 = new MyRDD(sc, 2, Nil) + val rddP2 = new MyRDD(sc, 2, List(new ShuffleDependency(rddR2, new HashPartitioner(2))), + tracker = mapOutputTracker) + val psdP2 = new PipelinedShuffleDependency(rddP2, new HashPartitioner(2)) + val rddC = new MyRDD(sc, 2, List(psdP1, psdP2), tracker = mapOutputTracker) + submit(rddC, Array(0, 1)) + + // Initially only the two regular roots R1, R2 run; P1, P2, C all wait. + assert(taskSets.size === 2, s"expected R1, R2 submitted, got ${taskSets.size}") + assert(scheduler.waitingStages.exists(_.rdd eq rddC)) + + // R1 completes -> P1 runs and reconsiders C. But P2 is still parked, so C must re-park. + val idR1 = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddR1).get.stageId + completeShuffleMapStageSuccessfully(idR1, 0, 2) + assert(scheduler.runningStages.exists(_.rdd eq rddP1), "P1 should be running") + assert(!scheduler.runningStages.exists(_.rdd eq rddP2), "P2 should not be running yet") + assert(scheduler.waitingStages.exists(_.rdd eq rddC), + "C must re-park while its second pipelined parent P2 is not yet running") + + // R2 completes -> P2 runs and reconsiders C; now both pipelined parents run, so C co-schedules. + val idR2 = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddR2).get.stageId + completeShuffleMapStageSuccessfully(idR2, 0, 2) + assert(scheduler.runningStages.exists(_.rdd eq rddC), "C should now be co-scheduled") + + val idP1 = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddP1).get.stageId + completeShuffleMapStageSuccessfully(idP1, 0, 2) + val idP2 = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddP2).get.stageId + completeShuffleMapStageSuccessfully(idP2, 0, 2) + val tsC = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddC).get + complete(tsC, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: one producer feeding two pipelined consumers reconsiders both") { + // R --regular--> P --pipelined--> C1 and C2 ; C1,C2 --regular--> D. + // When P starts (R done), reconsideration must resubmit BOTH parked consumers C1 and C2. + // (Guards the `.filter` in submitWaitingPipelinedChildStages against a `.find` regression.) + val rddR = new MyRDD(sc, 2, Nil) + val rddP = new MyRDD(sc, 2, List(new ShuffleDependency(rddR, new HashPartitioner(2))), + tracker = mapOutputTracker) + val rddC1 = new MyRDD(sc, 2, List(new PipelinedShuffleDependency(rddP, new HashPartitioner(2))), + tracker = mapOutputTracker) + val rddC2 = new MyRDD(sc, 2, List(new PipelinedShuffleDependency(rddP, new HashPartitioner(2))), + tracker = mapOutputTracker) + val rddD = new MyRDD(sc, 2, + List(new ShuffleDependency(rddC1, new HashPartitioner(2)), + new ShuffleDependency(rddC2, new HashPartitioner(2))), + tracker = mapOutputTracker) + submit(rddD, Array(0, 1)) + + // Only R runs initially; P, C1, C2 wait (D waits on its regular parents C1, C2). + assert(taskSets.size === 1, s"expected only R submitted, got ${taskSets.size}") + + // R completes -> P runs and reconsiders BOTH C1 and C2. + val idR = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddR).get.stageId + completeShuffleMapStageSuccessfully(idR, 0, 2) + assert(scheduler.runningStages.exists(_.rdd eq rddC1), "C1 should be co-scheduled with P") + assert(scheduler.runningStages.exists(_.rdd eq rddC2), "C2 should be co-scheduled with P") + + // Drain P, then C1, C2, then D. + val idP = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddP).get.stageId + completeShuffleMapStageSuccessfully(idP, 0, 2) + val idC1 = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddC1).get.stageId + completeShuffleMapStageSuccessfully(idC1, 0, 2) + val idC2 = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddC2).get.stageId + completeShuffleMapStageSuccessfully(idC2, 0, 2) + val tsD = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddD).get + complete(tsD, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: reconsideration cascades transitively (B->C->D behind a regular root)") { + // Z --regular--> B --pipelined--> C --pipelined--> D. When Z completes, B runs and reconsiders + // C; co-scheduling C must in turn reconsider D (transitive cascade through the co-schedule + // branch's submitWaitingPipelinedChildStages call). + val rddZ = new MyRDD(sc, 2, Nil) + val rddB = new MyRDD(sc, 2, List(new ShuffleDependency(rddZ, new HashPartitioner(2))), + tracker = mapOutputTracker) + val rddC = new MyRDD(sc, 2, List(new PipelinedShuffleDependency(rddB, new HashPartitioner(2))), + tracker = mapOutputTracker) + val rddD = new MyRDD(sc, 2, List(new PipelinedShuffleDependency(rddC, new HashPartitioner(2))), + tracker = mapOutputTracker) + submit(rddD, Array(0, 1)) + + assert(taskSets.size === 1, s"expected only Z submitted, got ${taskSets.size}") + + // Z completes -> B runs -> reconsiders C -> co-schedules C -> reconsiders D -> co-schedules D. + val idZ = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddZ).get.stageId + completeShuffleMapStageSuccessfully(idZ, 0, 2) + assert(scheduler.runningStages.exists(_.rdd eq rddB), "B running") + assert(scheduler.runningStages.exists(_.rdd eq rddC), "C co-scheduled with B") + assert(scheduler.runningStages.exists(_.rdd eq rddD), "D co-scheduled transitively with C") + + val idB = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddB).get.stageId + completeShuffleMapStageSuccessfully(idB, 0, 2) + val idC = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddC).get.stageId + completeShuffleMapStageSuccessfully(idC, 0, 2) + val tsD = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddD).get + complete(tsD, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: reconsidered consumer is not submitted twice when producer completes") { + // R --regular--> B --pipelined--> C. C is co-scheduled when B starts (reconsideration). When B + // later COMPLETES, submitWaitingChildStages(B) must NOT submit C a second time. + val rddR = new MyRDD(sc, 2, Nil) + val rddB = new MyRDD(sc, 2, List(new ShuffleDependency(rddR, new HashPartitioner(2))), + tracker = mapOutputTracker) + val rddC = new MyRDD(sc, 2, List(new PipelinedShuffleDependency(rddB, new HashPartitioner(2))), + tracker = mapOutputTracker) + submit(rddC, Array(0, 1)) + + val idR = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddR).get.stageId + completeShuffleMapStageSuccessfully(idR, 0, 2) + // C co-scheduled exactly once via reconsideration. + def cTaskSets = taskSets.count(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddC) + assert(cTaskSets === 1, "C should be submitted once by reconsideration") + + // B completes -> submitWaitingChildStages(B) fires, but C is already running, not waiting. + val idB = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddB).get.stageId + completeShuffleMapStageSuccessfully(idB, 0, 2) + assert(cTaskSets === 1, "C must not be submitted a second time when B completes") + + val tsC = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddC).get + complete(tsC, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + 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. + sc.conf.set(config.SPECULATION_ENABLED, true) + try { + 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)) + // Not rejected: the producer is submitted and the job proceeds normally. + assert(taskSets.nonEmpty, "regular-shuffle job must not be rejected under speculation") + completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2) + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } finally { + sc.conf.set(config.SPECULATION_ENABLED, false) + } + } + + test("pipelined shuffle: speculation is rejected for a job with a pipelined dependency") { + // Speculation races a producer copy against a consumer reading its partial output; reject. + // The scheduler reads sc.conf live, so toggling it here takes effect for the submit below. + sc.conf.set(config.SPECULATION_ENABLED, true) + 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 = {} + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(consumerRdd, Array(0, 1), listener = failListener) + assert(failure.get() != null, + "job with pipelined dependency + speculation should be rejected") + assert(failure.get().getMessage.contains("Speculative execution is not supported")) + // No task sets were submitted for the rejected job. + assert(taskSets.isEmpty) + assertDataStructuresEmpty() + } finally { + sc.conf.set(config.SPECULATION_ENABLED, false) + } + } + + test("pipelined shuffle: submitting a pipelined dependency as a map-stage job is rejected") { + // A map-stage job (SparkContext.submitMapStage, e.g. AQE's ShuffleExchangeExec) materializes a + // shuffle to produce map-output statistics. A PipelinedShuffleDependency cannot serve that -- it + // is transient with no durable map output (getStatistics would return all-zero stats), and its + // producer/consumer co-scheduling makes speculation unsafe. It is rejected outright, regardless + // of speculation. No partial scheduler state is left behind. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + 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) + } + submitMapStage(pipelinedDep, listener = failListener) + assert(failure.get() != null, + "submitMapStage with a pipelined dependency must be rejected") + assert(failure.get().getMessage.contains("cannot be submitted as a map-stage job"), + s"expected a map-stage-rejection error, got: ${failure.get().getMessage}") + // No task sets were submitted for the rejected map-stage job, and no partial state remains. + assert(taskSets.isEmpty) + assertDataStructuresEmpty() + } + + test("regular shuffle map-stage job with speculation enabled is NOT rejected (inertness)") { + // The map-stage speculation fail-fast must be inert for a regular ShuffleDependency. + sc.conf.set(config.SPECULATION_ENABLED, true) + try { + val producerRdd = new MyRDD(sc, 2, Nil) + val regularDep = new ShuffleDependency(producerRdd, new HashPartitioner(2)) + submitMapStage(regularDep) + // Not rejected: the map stage is submitted and runs normally. + assert(taskSets.nonEmpty, "regular-shuffle map-stage job must not be rejected under speculation") + completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2) + assertDataStructuresEmpty() + } finally { + sc.conf.set(config.SPECULATION_ENABLED, false) + } + } } class DAGSchedulerAbortStageOffSuite extends DAGSchedulerSuite { From 0050318874af108ec5d82d6424fce44e97591a52 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sat, 11 Jul 2026 17:07:28 +0000 Subject: [PATCH 02/54] [SPARK-XXXXX][CORE] Fail fast when a pipelined group cannot fit in cluster capacity Best-effort admission check for pipelined groups. All member stages of a group must run concurrently, so if the group's total task demand exceeds the cluster's total concurrent-task capacity it can never be co-resident and, lacking an out-of-band slot reservation, would deadlock (the consumer holds slots waiting for producer output while the producer cannot get slots to produce). We fail fast with a clear CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT error instead. When a pipelined group is about to be co-scheduled, submitStage computes the group's full pipelined-connected component (pipelinedGroupOf) and compares its summed task demand against maxConcurrentTasksForStage (production: sc.maxNumConcurrentTasks for the stage's resource profile). If demand exceeds capacity, the job is aborted; the already-launched producers are torn down by the normal job-abort path (cancelRunningIndependentStages). This is deliberately best-effort, not the atomic gang reservation deferred to a later hardening step: it compares the whole group's demand against TOTAL capacity (not free slots), checked once at co-schedule time. That converts the common under-provisioned case (a group that can never fit) into a clear error rather than a hang; races against other concurrently admitting work are left to the future atomic version. Inert for jobs with no pipelined dependency. maxConcurrentTasksForStage is a protected seam so tests can control reported capacity without changing the cluster's core count. Tests (DAGSchedulerSuite): a group too large to co-fit fails fast with CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT; a group that fits is co-scheduled normally; producer/consumer task sets marked isPipelined; regular task sets are not. Co-authored-by: Isaac --- .../resources/error/error-conditions.json | 6 + .../apache/spark/scheduler/DAGScheduler.scala | 130 +++++++++++++++++- .../spark/scheduler/DAGSchedulerSuite.scala | 83 +++++++++++ 3 files changed, 213 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 875be3d751580..b17f8c5facad3 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -1126,6 +1126,12 @@ ], "sqlState" : "0A000" }, + "CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT" : { + "message" : [ + "Cannot run the pipelined stage group: it needs concurrent task slots to run all its stages together, but only are currently free. The stages of a pipelined group must run concurrently, so the cluster must have enough free slots for all of them at once. Provision more executors, or reduce the parallelism of the query." + ], + "sqlState" : "53000" + }, "CONCURRENT_STREAM_LOG_UPDATE" : { "message" : [ "Concurrent update to the log. Multiple streaming jobs detected for .", 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 293cd73b4e383..2fb921fb36588 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -1426,6 +1426,98 @@ private[spark] class DAGScheduler( case _ => false } + /** + * The full pipelined group `stage` belongs to: the connected component of the stage graph over + * pipelined edges (both the producers `stage` reads through a [[PipelinedShuffleDependency]] and, + * transitively, their pipelined producers/consumers). Walks `parents` and the shuffle-map stages + * of pipelined dependencies. Returns just `stage` if it has no pipelined edge. + */ + private def pipelinedGroupOf(stage: Stage): Set[Stage] = { + val group = new HashSet[Stage] + val toVisit = new ListBuffer[Stage] + toVisit += stage + while (toVisit.nonEmpty) { + val s = toVisit.remove(0) + if (group.add(s)) { + // Pipelined producers this stage reads (direct pipelined parent shuffle-map stages). + s.parents.foreach { + case m: ShuffleMapStage + if m.shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]] => toVisit += m + case _ => + } + // Pipelined consumers/producers reachable the other direction: any known stage co-scheduled + // with s via a pipelined edge. We only have parent links, so also pull in stages whose + // pipelined parent is s (search the created stages). + stageIdToStage.valuesIterator.foreach { other => + if (!group.contains(other) && other.parents.contains(s) && + isPipelinedProducer(s)) { + toVisit += other + } + } + } + } + group.toSet + } + + /** + * Best-effort admission check for a pipelined group: all its member stages must run concurrently, + * so the cluster must have enough task slots for the whole group at once. If the group's total + * demand exceeds the cluster's total concurrent-task capacity it can never be co-resident and, + * lacking an out-of-band slot reservation, would deadlock -- so we fail fast with a clear error. + * + * This is deliberately best-effort, not the atomic gang reservation deferred to a later hardening + * step: it compares the whole group's demand against TOTAL capacity (not free slots), checked once + * when the group is first co-scheduled. That catches the common under-provisioned case (a group + * that can never fit) as a clear error rather than a hang; races against other concurrently + * admitting work are left to the future atomic version. + * + * Returns Some((demand, totalSlots)) when the group does not fit, else None. + */ + private def pipelinedGroupExceedsCapacity(stage: Stage): Option[(Int, Int)] = { + val group = pipelinedGroupOf(stage) + val demand = group.toSeq.map(_.numTasks).sum + val totalSlots = maxConcurrentTasksForStage(stage) + if (demand > totalSlots) Some((demand, totalSlots)) else None + } + + /** + * The cluster's total concurrent-task capacity for `stage`'s resource profile. Extracted as a + * seam so tests can control it without changing the cluster's core count. Production reads it + * from the scheduler backend, exactly as barrier's slot check does. + */ + protected def maxConcurrentTasksForStage(stage: Stage): Int = { + val rp = sc.resourceProfileManager.resourceProfileFromId(stage.resourceProfileId) + sc.maxNumConcurrentTasks(rp) + } + + /** + * Whether `stage` is a member of a pipelined group -- i.e. it is connected to another stage by a + * [[PipelinedShuffleDependency]], either as the producer (it writes such a shuffle) or as a + * consumer (one of its direct parent shuffle dependencies is pipelined). Members run + * concurrently over a transient shuffle that cannot be re-read in isolation, so a member's task + * failure must fail the whole group rather than resubmit one stage; the task scheduler keys its + * fail-fast behavior off this (via TaskSet.isPipelined). False for any stage in a job with no + * pipelined dependency. + */ + 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 + } + } + /** Finds the earliest-created active job that needs the stage */ // TODO: Probably should actually find among the active jobs that need this // stage the one with the highest priority (highest-priority pool, earliest created). @@ -1743,6 +1835,18 @@ private[spark] class DAGScheduler( submitStage(parent) } + // Submitting a parent can abort the job during this recursion: for a pipelined chain + // A->B->C, submitStage(C) recurses into submitStage(B) whose group slot-check covers the + // whole component {A,B,C} and may abortStage(B), which cleans up A/B/C (including this + // stage) and fails the job. If that happened, `stage` is no longer registered; do not + // proceed to co-schedule or re-park it (re-parking would re-insert a job-less stage into + // waitingStages -- a scheduler-state leak). Inert for a non-aborting submit. + if (!stageIdToStage.contains(stage.id)) { + logInfo(log"${MDC(STAGE, stage)} was removed during parent submission (its job was " + + log"aborted); not co-scheduling or re-parking it") + return + } + // A missing parent reached through a PipelinedShuffleDependency ("pipelined parent") // is incrementally readable: this stage may run before that parent materializes, so // the two are co-scheduled. `missing` already holds the direct parent shuffle-map @@ -1761,12 +1865,26 @@ private[spark] class DAGScheduler( pipelinedMissing.nonEmpty && pipelinedMissing.forall(runningStages.contains) if (regularMissing.isEmpty && allPipelinedParentsRunning) { - logInfo(log"Submitting ${MDC(STAGE, stage)} concurrently with its running pipelined " + - log"producer(s) ${MDC(MISSING_PARENT_STAGES, pipelinedMissing)}") - submitMissingTasks(stage, jobId.get) - // This stage is now running; if it is itself the pipelined producer of a waiting - // consumer, co-schedule that consumer too. - submitWaitingPipelinedChildStages(stage) + // Best-effort slot check: the whole pipelined group must run at once, so it must fit + // in the cluster's total concurrent-task capacity. If it cannot, fail fast with a + // clear error rather than deadlock (there is no out-of-band slot reservation in v1). + pipelinedGroupExceedsCapacity(stage) match { + case Some((demand, totalSlots)) => + abortStage(stage, s"Cannot co-schedule pipelined stage group: needs $demand " + + s"concurrent task slots but the cluster has only $totalSlots.", + Some(new SparkException( + errorClass = "CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT", + messageParameters = scala.collection.immutable.Map( + "numTasks" -> demand.toString, "numSlots" -> totalSlots.toString), + cause = null))) + case None => + logInfo(log"Submitting ${MDC(STAGE, stage)} concurrently with its running " + + log"pipelined producer(s) ${MDC(MISSING_PARENT_STAGES, pipelinedMissing)}") + submitMissingTasks(stage, jobId.get) + // This stage is now running; if it is itself the pipelined producer of a waiting + // consumer, co-schedule that consumer too. + submitWaitingPipelinedChildStages(stage) + } } else { waitingStages += stage } 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 a58e60de0325c..4d1e36456d175 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -380,6 +380,13 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti shuffleMergeRegister: Boolean = true ) extends DAGScheduler( sc, taskScheduler, listenerBus, mapOutputTracker, blockManagerMaster, env, clock) { + + // Capacity reported to the pipelined-group slot check. Defaults high so scheduling-logic tests + // are not gated by the real (local[2]) backend capacity; a test can lower it to exercise the + // fail-fast path. + @volatile var maxConcurrentTasksForTest: Int = 1000 + override protected def maxConcurrentTasksForStage(stage: Stage): Int = maxConcurrentTasksForTest + /** * Schedules shuffle merge finalize. */ @@ -6545,6 +6552,82 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti sc.conf.set(config.SPECULATION_ENABLED, false) } } + + // ========================================================================================== + // Gang admission / slot check (spec S4, S4.1) + // ========================================================================================== + + test("pipelined shuffle: a group too large to co-fit fails fast with INSUFFICIENT_SLOT") { + // Constrain reported capacity to 3 slots. A pipelined group of producer(2) + consumer(2) = 4 + // tasks exceeds it and can never be co-resident, so it must fail fast rather than deadlock. + val producerRdd = new MyRDD(sc, 2, Nil) // touch sc first so `scheduler` is initialized + scheduler.asInstanceOf[MyDAGScheduler].maxConcurrentTasksForTest = 3 + try { + 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, "an over-large pipelined group must fail the job") + assert(failure.get().getMessage.contains("CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT") || + failure.get().getMessage.contains("concurrent task slots"), + s"expected an insufficient-slot error, got: ${failure.get().getMessage}") + assertDataStructuresEmpty() + } finally { + scheduler.asInstanceOf[MyDAGScheduler].maxConcurrentTasksForTest = 1000 + } + } + + test("pipelined shuffle: an over-capacity 3-stage chain aborts without leaking a stage in " + + "waitingStages (nested-frame abort)") { + // Regression for a nested-frame abort: for a pipelined chain A->B->C, submitStage(C) recurses + // into submitStage(B), whose group slot-check covers the whole component {A,B,C} and aborts the + // job from that nested frame (cleaning up A/B/C). Control then unwinds to the outer + // submitStage(C) frame; without the "was I removed during recursion?" guard, its else branch + // would re-park the already-removed C into waitingStages -- a job-less scheduler-state leak. The + // 2-stage over-capacity test does not exercise this (there the abort fires in the consumer's own + // frame). Assert the abort leaves NO stage behind (assertDataStructuresEmpty checks waitingStages). + val rddA = new MyRDD(sc, 2, Nil) // touch sc first so `scheduler` is initialized + scheduler.asInstanceOf[MyDAGScheduler].maxConcurrentTasksForTest = 3 + try { + val depAB = new PipelinedShuffleDependency(rddA, new HashPartitioner(2)) + val rddB = new MyRDD(sc, 2, List(depAB), tracker = mapOutputTracker) + val depBC = new PipelinedShuffleDependency(rddB, new HashPartitioner(2)) + val rddC = new MyRDD(sc, 2, List(depBC), 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(rddC, Array(0, 1), listener = failListener) + + assert(failure.get() != null, "an over-large 3-stage pipelined chain must fail the job") + assert(failure.get().getMessage.contains("CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT") || + failure.get().getMessage.contains("concurrent task slots"), + s"expected an insufficient-slot error, got: ${failure.get().getMessage}") + // The key assertion: no stage (esp. the outer consumer C) is left behind in waitingStages. + assertDataStructuresEmpty() + } finally { + scheduler.asInstanceOf[MyDAGScheduler].maxConcurrentTasksForTest = 1000 + } + } + + test("pipelined shuffle: a group that fits is co-scheduled (slot check passes)") { + // Producer (2) + consumer (2) = 4 tasks <= 16 slots: co-scheduled normally, no abort. + 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, "group fits, so producer and consumer are co-scheduled") + completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2) + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } } class DAGSchedulerAbortStageOffSuite extends DAGSchedulerSuite { From b2c9d9e836954a7fd602fefd1c4c217d680df89e Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Wed, 15 Jul 2026 00:17:53 +0000 Subject: [PATCH 03/54] [SPARK-XXXXX][CORE] Admit pipelined groups against free slots (spec S4.1) Upgrade the gang-admission slot check from total capacity to currently-FREE slots: free = per-profile total capacity minus tasks already running for OTHER work (other groups / regular jobs), excluding the group's own already-running members. Adds TaskSchedulerImpl.runningTasksForOtherWorkInProfile (single-lock, resource-profile-scoped, zombie-filtered) and a DAGScheduler seam. Comparing against total capacity would admit a group that fits in principle but cannot co-fit beside a busy neighbor, then hang; free-slot admission fails fast instead. Also bumps the incremental-manager config .version to 4.3.0. --- .../apache/spark/scheduler/DAGScheduler.scala | 62 +++++++++---- .../spark/scheduler/TaskSchedulerImpl.scala | 24 +++++ .../spark/scheduler/DAGSchedulerSuite.scala | 87 +++++++++++++++++-- .../scheduler/TaskSchedulerImplSuite.scala | 29 +++++++ 4 files changed, 179 insertions(+), 23 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 2fb921fb36588..a31c1f9bfc8d7 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -1461,23 +1461,39 @@ private[spark] class DAGScheduler( /** * Best-effort admission check for a pipelined group: all its member stages must run concurrently, - * so the cluster must have enough task slots for the whole group at once. If the group's total - * demand exceeds the cluster's total concurrent-task capacity it can never be co-resident and, - * lacking an out-of-band slot reservation, would deadlock -- so we fail fast with a clear error. + * so the cluster must have enough task slots for the whole group at once. If the group's demand + * exceeds what is currently available it can never be co-resident and, lacking an out-of-band slot + * reservation, would deadlock -- so we fail fast with a clear error. * - * This is deliberately best-effort, not the atomic gang reservation deferred to a later hardening - * step: it compares the whole group's demand against TOTAL capacity (not free slots), checked once - * when the group is first co-scheduled. That catches the common under-provisioned case (a group - * that can never fit) as a clear error rather than a hang; races against other concurrently - * admitting work are left to the future atomic version. + * Demand is compared against FREE slots, not total capacity (spec S4.1): free = total capacity + * minus the tasks already running for OTHER work (other groups and regular jobs) in the SAME + * resource profile. Comparing against total capacity would admit a group that fits the cluster in + * principle but cannot co-fit right now beside a busy neighbor, then hang. Occupancy excludes the + * group's OWN already-running members (e.g. a producer submitted just before its consumer is + * admitted) -- otherwise the group would be charged for its own slots and could fail a check it + * actually fits. * - * Returns Some((demand, totalSlots)) when the group does not fit, else None. + * Occupancy is resource-profile-scoped to match `totalSlots` (which is per-profile): counting + * other profiles' running tasks against one profile's capacity would spuriously reject a fitting + * group whenever an unrelated profile is busy. v1 requires a group to be single-profile (spec S9), + * so the whole group shares `stage`'s profile. + * + * This is best-effort, checked once when the group is first co-scheduled; it is NOT the atomic + * gang reservation deferred to a later hardening step. Without that reservation a race between two + * groups admitting concurrently can still transiently over-admit (each sees the other's slots as + * free before either's tasks register as running); the reservation closes that race later. What + * this does guarantee now is that a group is never admitted against slots a busy neighbor already + * holds, which is the common single-group-on-a-busy-cluster hang. + * + * Returns Some((demand, freeSlots)) when the group does not fit, else None. */ private def pipelinedGroupExceedsCapacity(stage: Stage): Option[(Int, Int)] = { val group = pipelinedGroupOf(stage) val demand = group.toSeq.map(_.numTasks).sum val totalSlots = maxConcurrentTasksForStage(stage) - if (demand > totalSlots) Some((demand, totalSlots)) else None + val occupiedByOthers = runningTasksForOtherWork(stage, group) + val freeSlots = math.max(0, totalSlots - occupiedByOthers) + if (demand > freeSlots) Some((demand, freeSlots)) else None } /** @@ -1490,6 +1506,19 @@ private[spark] class DAGScheduler( sc.maxNumConcurrentTasks(rp) } + /** + * Tasks currently running, in `stage`'s resource profile, for work OTHER than the given pipelined + * `group` (whose own already-running members must not be charged against the group's admission). + * Resource-profile-scoped to match `maxConcurrentTasksForStage`. Extracted as a seam so tests can + * control occupancy without launching real tasks; returns 0 for a non-TaskSchedulerImpl backend. + */ + protected def runningTasksForOtherWork(stage: Stage, group: Set[Stage]): Int = + taskScheduler match { + case impl: TaskSchedulerImpl => + impl.runningTasksForOtherWorkInProfile(stage.resourceProfileId, group.map(_.id)) + case _ => 0 + } + /** * Whether `stage` is a member of a pipelined group -- i.e. it is connected to another stage by a * [[PipelinedShuffleDependency]], either as the producer (it writes such a shuffle) or as a @@ -1865,17 +1894,18 @@ private[spark] class DAGScheduler( pipelinedMissing.nonEmpty && pipelinedMissing.forall(runningStages.contains) if (regularMissing.isEmpty && allPipelinedParentsRunning) { - // Best-effort slot check: the whole pipelined group must run at once, so it must fit - // in the cluster's total concurrent-task capacity. If it cannot, fail fast with a - // clear error rather than deadlock (there is no out-of-band slot reservation in v1). + // Best-effort slot check: the whole pipelined group must run at once, so its demand + // must fit in the currently-FREE slots (total capacity minus what other work is + // already running; spec S4.1). If it cannot, fail fast with a clear error rather than + // deadlock (there is no out-of-band slot reservation in v1). pipelinedGroupExceedsCapacity(stage) match { - case Some((demand, totalSlots)) => + case Some((demand, freeSlots)) => abortStage(stage, s"Cannot co-schedule pipelined stage group: needs $demand " + - s"concurrent task slots but the cluster has only $totalSlots.", + s"concurrent task slots but only $freeSlots are currently free.", Some(new SparkException( errorClass = "CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT", messageParameters = scala.collection.immutable.Map( - "numTasks" -> demand.toString, "numSlots" -> totalSlots.toString), + "numTasks" -> demand.toString, "numSlots" -> freeSlots.toString), cause = null))) case None => logInfo(log"Submitting ${MDC(STAGE, stage)} concurrently with its running " + diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala index 618c8eb459026..d1669343a6a93 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala @@ -160,6 +160,30 @@ private[spark] class TaskSchedulerImpl( executorIdToRunningTaskIds.toMap.transform((_, v) => v.size) } + /** + * The number of tasks currently running for the given resource profile that belong to work + * OTHER than the stages in `excludeStageIds`. Used by the pipelined-group free-slot admission + * check (see DAGScheduler): it must compare the group's demand against the slots left free by + * everything else running in the SAME resource profile, so it excludes the group's own already- + * running members (whose tasks would otherwise be charged against the group's own admission). + * + * Computed from the TaskSetManagers so it is resource-profile-scoped and consistently filters + * zombie attempts, and taken under a single lock so the count is one consistent snapshot (both + * the per-profile total and the excluded members are read together). + */ + def runningTasksForOtherWorkInProfile( + resourceProfileId: Int, excludeStageIds: Set[Int]): Int = synchronized { + taskSetsByStageIdAndAttempt.iterator.flatMap { case (stageId, attempts) => + if (excludeStageIds.contains(stageId)) { + Iterator.empty + } else { + attempts.valuesIterator + .filter(tsm => !tsm.isZombie && tsm.taskSet.resourceProfileId == resourceProfileId) + .map(_.runningTasks) + } + }.sum + } + // The set of executors we have on each host; this is used to compute hostsAlive, which // in turn is used to decide when we can attain data locality on a given host protected val hostToExecutors = new HashMap[String, HashSet[String]] 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 4d1e36456d175..26950eb7e0367 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -387,6 +387,12 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti @volatile var maxConcurrentTasksForTest: Int = 1000 override protected def maxConcurrentTasksForStage(stage: Stage): Int = maxConcurrentTasksForTest + // Seam for the free-slot admission check (spec S4.1): occupancy by OTHER work. Default 0 so + // existing tests see a fully-free cluster; a test can set it to model a busy neighbor. + @volatile var runningTasksForOtherWorkForTest: (Stage, Set[Stage]) => Int = (_, _) => 0 + override protected def runningTasksForOtherWork(stage: Stage, group: Set[Stage]): Int = + runningTasksForOtherWorkForTest(stage, group) + /** * Schedules shuffle merge finalize. */ @@ -6271,7 +6277,8 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } - test("pipelined shuffle: a pipelined producer behind a regular shuffle is NOT co-scheduled early") { + test("pipelined shuffle: a pipelined producer behind a regular shuffle is NOT co-scheduled " + + "early") { // A --regular--> B --pipelined--> C. B cannot run until A materializes, so C must NOT be // co-scheduled with B while B is still parked. (Regression: the naive check co-scheduled C // against a not-yet-running B, stranding C -- a hang under slot pressure.) @@ -6455,7 +6462,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti val idR = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddR).get.stageId completeShuffleMapStageSuccessfully(idR, 0, 2) // C co-scheduled exactly once via reconsideration. - def cTaskSets = taskSets.count(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddC) + def cTaskSets: Int = taskSets.count(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddC) assert(cTaskSets === 1, "C should be submitted once by reconsideration") // B completes -> submitWaitingChildStages(B) fires, but C is already running, not waiting. @@ -6516,7 +6523,8 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti test("pipelined shuffle: submitting a pipelined dependency as a map-stage job is rejected") { // A map-stage job (SparkContext.submitMapStage, e.g. AQE's ShuffleExchangeExec) materializes a - // shuffle to produce map-output statistics. A PipelinedShuffleDependency cannot serve that -- it + // shuffle to produce map-output statistics. A PipelinedShuffleDependency cannot serve that -- + // it // is transient with no durable map output (getStatistics would return all-zero stats), and its // producer/consumer co-scheduling makes speculation unsafe. It is rejected outright, regardless // of speculation. No partial scheduler state is left behind. @@ -6545,7 +6553,8 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti val regularDep = new ShuffleDependency(producerRdd, new HashPartitioner(2)) submitMapStage(regularDep) // Not rejected: the map stage is submitted and runs normally. - assert(taskSets.nonEmpty, "regular-shuffle map-stage job must not be rejected under speculation") + assert(taskSets.nonEmpty, + "regular-shuffle map-stage job must not be rejected under speculation") completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2) assertDataStructuresEmpty() } finally { @@ -6588,9 +6597,12 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti // into submitStage(B), whose group slot-check covers the whole component {A,B,C} and aborts the // job from that nested frame (cleaning up A/B/C). Control then unwinds to the outer // submitStage(C) frame; without the "was I removed during recursion?" guard, its else branch - // would re-park the already-removed C into waitingStages -- a job-less scheduler-state leak. The - // 2-stage over-capacity test does not exercise this (there the abort fires in the consumer's own - // frame). Assert the abort leaves NO stage behind (assertDataStructuresEmpty checks waitingStages). + // would re-park the already-removed C into waitingStages -- a job-less scheduler-state leak. + // The + // 2-stage over-capacity test does not exercise this (there the abort fires in the consumer's + // own + // frame). Assert the abort leaves NO stage behind (assertDataStructuresEmpty checks + // waitingStages). val rddA = new MyRDD(sc, 2, Nil) // touch sc first so `scheduler` is initialized scheduler.asInstanceOf[MyDAGScheduler].maxConcurrentTasksForTest = 3 try { @@ -6628,6 +6640,67 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assert(results === Map(0 -> 42, 1 -> 43)) assertDataStructuresEmpty() } + + test("pipelined shuffle: a group that fits total capacity but not FREE slots fails fast (S4.1)") { + // Spec S4.1: admission is decided against currently-FREE slots, not total capacity. A group of + // producer(2) + consumer(2) = 4 tasks fits a 10-slot cluster in principle, but if 8 slots are + // already occupied by OTHER work only 2 are free -- the group cannot co-fit right now and must + // fail fast rather than queue forever. Under the old total-capacity check this would wrongly be + // admitted (4 <= 10) and then hang. runningTasksForOtherWork already excludes the group's own + // members, so we report 8 directly. + val producerRdd = new MyRDD(sc, 2, Nil) // touch sc first so `scheduler` is initialized + val myScheduler = scheduler.asInstanceOf[MyDAGScheduler] + myScheduler.maxConcurrentTasksForTest = 10 + // 8 of 10 slots busy elsewhere -> 2 free + myScheduler.runningTasksForOtherWorkForTest = (_, _) => 8 + try { + 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 group that fits total but not free slots must fail the job (S4.1 free-slot check)") + assert(failure.get().getMessage.contains("CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT") || + failure.get().getMessage.contains("currently free"), + s"expected a free-slot insufficient-slot error, got: ${failure.get().getMessage}") + assertDataStructuresEmpty() + } finally { + myScheduler.maxConcurrentTasksForTest = 1000 + myScheduler.runningTasksForOtherWorkForTest = (_, _) => 0 + } + } + + test("pipelined shuffle: the group's own running members are not charged against its admission") { + // Occupancy for the free-slot check counts only OTHER work: the group's own already-running + // producer must not be charged, or a group that exactly fills free capacity would be wrongly + // rejected. Model a full cluster where the ONLY running tasks are this group's own: 4 total + // slots, other-work occupancy = 0 (the seam excludes the group) -> free = 4 >= demand 4 -> + // admitted. The exclusion itself is unit-tested against TaskSchedulerImpl separately; here we + // assert the DAGScheduler admits a demand==capacity group when nothing else is running. + val producerRdd = new MyRDD(sc, 2, Nil) + val myScheduler = scheduler.asInstanceOf[MyDAGScheduler] + myScheduler.maxConcurrentTasksForTest = 4 + myScheduler.runningTasksForOtherWorkForTest = (_, _) => 0 // only the group's own work runs + try { + 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, + "a group whose demand equals free capacity must be co-scheduled, not rejected") + completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2) + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } finally { + myScheduler.maxConcurrentTasksForTest = 1000 + myScheduler.runningTasksForOtherWorkForTest = (_, _) => 0 + } + } } class DAGSchedulerAbortStageOffSuite extends DAGSchedulerSuite { diff --git a/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala index c9c9e529405ec..c62f1de2798d6 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala @@ -2752,4 +2752,33 @@ class TaskSchedulerImplSuite extends SparkFunSuite with LocalSparkContext assert(!tsm.isInstanceOf[StructuredStreamingIdAwareSchedulerLogging]) } + test("runningTasksForOtherWorkInProfile excludes given stages and is resource-profile-scoped") { + // Backs the pipelined-group free-slot admission check (DAGScheduler): it must report tasks + // running for OTHER work in a given resource profile, excluding the group's own member stages. + val taskScheduler = setupScheduler() + val workerOffers = IndexedSeq( + new WorkerOffer("executor0", "host0", 4), + new WorkerOffer("executor1", "host1", 4)) + + // Stage 0 and stage 1 both in the DEFAULT profile; launch 2 tasks of each so both have running + // tasks. Distinct stageIds let us exclude one and count the other. + taskScheduler.submitTasks(FakeTask.createTaskSet(2, stageId = 0, stageAttemptId = 0)) + taskScheduler.submitTasks(FakeTask.createTaskSet(2, stageId = 1, stageAttemptId = 0)) + val launched = taskScheduler.resourceOffers(workerOffers).flatten + assert(launched.length === 4, "all four tasks should launch (4 cores per executor)") + + val defaultRp = ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID + // Nothing excluded: both stages' 4 tasks count. + assert(taskScheduler.runningTasksForOtherWorkInProfile(defaultRp, Set.empty) === 4) + // Exclude stage 0 (as if it were the group's own member): only stage 1's 2 tasks remain. + assert(taskScheduler.runningTasksForOtherWorkInProfile(defaultRp, Set(0)) === 2) + // Exclude both: zero "other" tasks. + assert(taskScheduler.runningTasksForOtherWorkInProfile(defaultRp, Set(0, 1)) === 0) + // A DIFFERENT (non-existent here) resource profile has no running tasks -- other-profile load + // must not be charged against this profile (finding A: RP-scoping). + val otherRpId = defaultRp + 12345 + assert(taskScheduler.runningTasksForOtherWorkInProfile(otherRpId, Set.empty) === 0, + "tasks in the default profile must not be counted against a different profile") + } + } From e85b735dd1cbcd5bfc078a5f8f5be81ce3c56698 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Fri, 17 Jul 2026 20:36:54 +0000 Subject: [PATCH 04/54] [SPARK-XXXXX][CORE] Pipelined group admission: count enqueued demand, add a slot-check disable flag Two updates to the pipelined-group slot admission check (spec S4.1), per an updated spec: - Count outstanding demand, not just running tasks. The occupancy of OTHER work in a resource profile now sums each task set's not-yet-completed tasks (numTasks - tasksSuccessful), i.e. running plus enqueued, instead of only the tasks actively running. A neighbor's queued backlog is real committed demand that can starve a co-scheduled group, so charging only running tasks could admit a group that then hangs once the backlog launches. Renamed TaskSchedulerImpl's helper to outstandingTasksForOtherWorkInProfile and the DAGScheduler seam to outstandingTasksForOtherWork to reflect the new semantics. - Add spark.scheduler.pipelinedGroup.slotCheck.enabled (internal, default true). When false, pipelinedGroupExceedsCapacity is skipped entirely and a group is co-scheduled unconditionally. This mirrors ConcurrentStageDAGScheduler's ability to disable its slot check, for deployments that admit capacity out-of-band (e.g. a slot reservation) and own admission themselves. Tests: the TaskSchedulerImpl unit test now submits more tasks than slots and asserts the count includes the enqueued tasks; a new DAGScheduler test asserts that with the check disabled an over-capacity group is co-scheduled rather than failed. Co-authored-by: Isaac --- .../spark/internal/config/package.scala | 15 ++++++ .../apache/spark/scheduler/DAGScheduler.scala | 31 ++++++++---- .../spark/scheduler/TaskSchedulerImpl.scala | 24 ++++----- .../spark/scheduler/DAGSchedulerSuite.scala | 49 ++++++++++++++----- .../scheduler/TaskSchedulerImplSuite.scala | 41 +++++++++------- 5 files changed, 110 insertions(+), 50 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/internal/config/package.scala b/core/src/main/scala/org/apache/spark/internal/config/package.scala index c8bd6410be27b..2f20efbc0a0ac 100644 --- a/core/src/main/scala/org/apache/spark/internal/config/package.scala +++ b/core/src/main/scala/org/apache/spark/internal/config/package.scala @@ -2006,6 +2006,21 @@ package object config { .checkValue(v => v > 0, "The max failures should be a positive value.") .createWithDefault(40) + private[spark] val PIPELINED_GROUP_SLOT_CHECK_ENABLED = + ConfigBuilder("spark.scheduler.pipelinedGroup.slotCheck.enabled") + .internal() + .doc("When true, before co-scheduling a pipelined-shuffle stage group the DAGScheduler " + + "checks that the group's total task demand fits in the currently free slots of its " + + "resource profile (total capacity minus the outstanding -- running plus enqueued -- " + + "tasks of other work), and fails the job with CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT " + + "rather than co-scheduling a group that cannot fit and deadlocking. Set to false for " + + "deployments that admit capacity out-of-band (e.g. a slot reservation), which then own " + + "admission. Only applies to jobs that use a pipelined shuffle dependency.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .booleanConf + .createWithDefault(true) + private[spark] val NUM_CANCELLED_JOB_GROUPS_TO_TRACK = ConfigBuilder("spark.scheduler.numCancelledJobGroupsToTrack") .doc("The maximum number of tracked job groups that are cancelled with " + 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 a31c1f9bfc8d7..89df4a5bad888 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -1482,17 +1482,25 @@ private[spark] class DAGScheduler( * gang reservation deferred to a later hardening step. Without that reservation a race between two * groups admitting concurrently can still transiently over-admit (each sees the other's slots as * free before either's tasks register as running); the reservation closes that race later. What - * this does guarantee now is that a group is never admitted against slots a busy neighbor already - * holds, which is the common single-group-on-a-busy-cluster hang. + * this does guarantee now is that a group is never admitted against slots a busy neighbor is + * already committed to (running or enqueued), which is the common single-group-on-a-busy-cluster + * hang. + * + * The check can be turned off with `spark.scheduler.pipelinedGroup.slotCheck.enabled=false` (for + * deployments that admit capacity out-of-band, e.g. via a slot reservation), in which case this + * always returns None and the group is co-scheduled unconditionally. * * Returns Some((demand, freeSlots)) when the group does not fit, else None. */ private def pipelinedGroupExceedsCapacity(stage: Stage): Option[(Int, Int)] = { + if (!sc.conf.get(config.PIPELINED_GROUP_SLOT_CHECK_ENABLED)) { + return None + } val group = pipelinedGroupOf(stage) val demand = group.toSeq.map(_.numTasks).sum val totalSlots = maxConcurrentTasksForStage(stage) - val occupiedByOthers = runningTasksForOtherWork(stage, group) - val freeSlots = math.max(0, totalSlots - occupiedByOthers) + val outstandingForOthers = outstandingTasksForOtherWork(stage, group) + val freeSlots = math.max(0, totalSlots - outstandingForOthers) if (demand > freeSlots) Some((demand, freeSlots)) else None } @@ -1507,15 +1515,18 @@ private[spark] class DAGScheduler( } /** - * Tasks currently running, in `stage`'s resource profile, for work OTHER than the given pipelined - * `group` (whose own already-running members must not be charged against the group's admission). - * Resource-profile-scoped to match `maxConcurrentTasksForStage`. Extracted as a seam so tests can - * control occupancy without launching real tasks; returns 0 for a non-TaskSchedulerImpl backend. + * Outstanding tasks (running plus enqueued), in `stage`'s resource profile, for work OTHER than + * the given pipelined `group` (whose own members must not be charged against the group's own + * admission). Counting enqueued tasks, not just running ones, means a busy neighbor's queued + * backlog is charged against capacity too, so a group is not admitted against slots that other + * work is already committed to using. Resource-profile-scoped to match + * `maxConcurrentTasksForStage`. Extracted as a seam so tests can control occupancy without + * launching real tasks; returns 0 for a non-TaskSchedulerImpl backend. */ - protected def runningTasksForOtherWork(stage: Stage, group: Set[Stage]): Int = + protected def outstandingTasksForOtherWork(stage: Stage, group: Set[Stage]): Int = taskScheduler match { case impl: TaskSchedulerImpl => - impl.runningTasksForOtherWorkInProfile(stage.resourceProfileId, group.map(_.id)) + impl.outstandingTasksForOtherWorkInProfile(stage.resourceProfileId, group.map(_.id)) case _ => 0 } diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala index d1669343a6a93..fba876717cef8 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala @@ -161,25 +161,27 @@ private[spark] class TaskSchedulerImpl( } /** - * The number of tasks currently running for the given resource profile that belong to work - * OTHER than the stages in `excludeStageIds`. Used by the pipelined-group free-slot admission - * check (see DAGScheduler): it must compare the group's demand against the slots left free by - * everything else running in the SAME resource profile, so it excludes the group's own already- - * running members (whose tasks would otherwise be charged against the group's own admission). + * The number of outstanding tasks for the given resource profile that belong to work OTHER than + * the stages in `excludeStageIds`. "Outstanding" is the not-yet-completed demand of each task set + * -- running plus enqueued (`numTasks - tasksSuccessful`) -- not just the tasks actively running, + * so a neighbor's queued backlog is charged against capacity too. Used by the pipelined-group + * slot admission check (see DAGScheduler): it compares the group's demand against the slots free + * by everything else in the SAME resource profile, so it excludes the group's own members (whose + * tasks would otherwise be charged against the group's own admission). * - * Computed from the TaskSetManagers so it is resource-profile-scoped and consistently filters - * zombie attempts, and taken under a single lock so the count is one consistent snapshot (both - * the per-profile total and the excluded members are read together). + * Computed from the TaskSetManagers so it is resource-profile-scoped, and taken under a single + * lock so the count is one consistent snapshot (both the per-profile total and the excluded + * members are read together). */ - def runningTasksForOtherWorkInProfile( + def outstandingTasksForOtherWorkInProfile( resourceProfileId: Int, excludeStageIds: Set[Int]): Int = synchronized { taskSetsByStageIdAndAttempt.iterator.flatMap { case (stageId, attempts) => if (excludeStageIds.contains(stageId)) { Iterator.empty } else { attempts.valuesIterator - .filter(tsm => !tsm.isZombie && tsm.taskSet.resourceProfileId == resourceProfileId) - .map(_.runningTasks) + .filter(_.taskSet.resourceProfileId == resourceProfileId) + .map(tsm => math.max(0, tsm.numTasks - tsm.tasksSuccessful)) } }.sum } 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 26950eb7e0367..2bf567d4c8c9d 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -387,11 +387,12 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti @volatile var maxConcurrentTasksForTest: Int = 1000 override protected def maxConcurrentTasksForStage(stage: Stage): Int = maxConcurrentTasksForTest - // Seam for the free-slot admission check (spec S4.1): occupancy by OTHER work. Default 0 so - // existing tests see a fully-free cluster; a test can set it to model a busy neighbor. - @volatile var runningTasksForOtherWorkForTest: (Stage, Set[Stage]) => Int = (_, _) => 0 - override protected def runningTasksForOtherWork(stage: Stage, group: Set[Stage]): Int = - runningTasksForOtherWorkForTest(stage, group) + // Seam for the free-slot admission check (spec S4.1): outstanding (running + enqueued) task + // demand of OTHER work. Default 0 so existing tests see a fully-free cluster; a test can set it + // to model a busy neighbor. + @volatile var outstandingTasksForOtherWorkForTest: (Stage, Set[Stage]) => Int = (_, _) => 0 + override protected def outstandingTasksForOtherWork(stage: Stage, group: Set[Stage]): Int = + outstandingTasksForOtherWorkForTest(stage, group) /** * Schedules shuffle merge finalize. @@ -6646,13 +6647,13 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti // producer(2) + consumer(2) = 4 tasks fits a 10-slot cluster in principle, but if 8 slots are // already occupied by OTHER work only 2 are free -- the group cannot co-fit right now and must // fail fast rather than queue forever. Under the old total-capacity check this would wrongly be - // admitted (4 <= 10) and then hang. runningTasksForOtherWork already excludes the group's own - // members, so we report 8 directly. + // admitted (4 <= 10) and then hang. outstandingTasksForOtherWork already excludes the group's + // own members, so we report 8 directly. val producerRdd = new MyRDD(sc, 2, Nil) // touch sc first so `scheduler` is initialized val myScheduler = scheduler.asInstanceOf[MyDAGScheduler] myScheduler.maxConcurrentTasksForTest = 10 // 8 of 10 slots busy elsewhere -> 2 free - myScheduler.runningTasksForOtherWorkForTest = (_, _) => 8 + myScheduler.outstandingTasksForOtherWorkForTest = (_, _) => 8 try { val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) @@ -6671,7 +6672,33 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } finally { myScheduler.maxConcurrentTasksForTest = 1000 - myScheduler.runningTasksForOtherWorkForTest = (_, _) => 0 + myScheduler.outstandingTasksForOtherWorkForTest = (_, _) => 0 + } + } + + test("pipelined shuffle: slot check can be disabled, admitting a group that exceeds capacity") { + // With spark.scheduler.pipelinedGroup.slotCheck.enabled=false the admission check is skipped + // entirely, so a group whose demand exceeds free slots is co-scheduled instead of failing fast. + // Deployments that admit capacity out-of-band (e.g. a slot reservation) rely on this. + val producerRdd = new MyRDD(sc, 2, Nil) + val myScheduler = scheduler.asInstanceOf[MyDAGScheduler] + myScheduler.maxConcurrentTasksForTest = 2 // demand 4 > 2 slots: would fail if the check ran + myScheduler.outstandingTasksForOtherWorkForTest = (_, _) => 0 + val prev = sc.conf.get(config.PIPELINED_GROUP_SLOT_CHECK_ENABLED) + sc.conf.set(config.PIPELINED_GROUP_SLOT_CHECK_ENABLED, false) + try { + 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, + "with the slot check disabled, the over-capacity group is co-scheduled, not rejected") + completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2) + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } finally { + myScheduler.maxConcurrentTasksForTest = 1000 + sc.conf.set(config.PIPELINED_GROUP_SLOT_CHECK_ENABLED, prev) } } @@ -6685,7 +6712,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti val producerRdd = new MyRDD(sc, 2, Nil) val myScheduler = scheduler.asInstanceOf[MyDAGScheduler] myScheduler.maxConcurrentTasksForTest = 4 - myScheduler.runningTasksForOtherWorkForTest = (_, _) => 0 // only the group's own work runs + myScheduler.outstandingTasksForOtherWorkForTest = (_, _) => 0 // only the group's own work runs try { val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) @@ -6698,7 +6725,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } finally { myScheduler.maxConcurrentTasksForTest = 1000 - myScheduler.runningTasksForOtherWorkForTest = (_, _) => 0 + myScheduler.outstandingTasksForOtherWorkForTest = (_, _) => 0 } } } diff --git a/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala index c62f1de2798d6..ef21af6096487 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala @@ -2752,32 +2752,37 @@ class TaskSchedulerImplSuite extends SparkFunSuite with LocalSparkContext assert(!tsm.isInstanceOf[StructuredStreamingIdAwareSchedulerLogging]) } - test("runningTasksForOtherWorkInProfile excludes given stages and is resource-profile-scoped") { - // Backs the pipelined-group free-slot admission check (DAGScheduler): it must report tasks - // running for OTHER work in a given resource profile, excluding the group's own member stages. + test("outstandingTasksForOtherWorkInProfile counts running + enqueued, excludes given stages, " + + "and is resource-profile-scoped") { + // Backs the pipelined-group slot admission check (DAGScheduler): it must report the OUTSTANDING + // (running + enqueued, i.e. numTasks - tasksSuccessful) task demand of OTHER work in a given + // resource profile, excluding the group's own member stages. val taskScheduler = setupScheduler() + // Only 3 total cores, but each stage has 4 tasks -> some tasks launch, the rest stay enqueued. val workerOffers = IndexedSeq( - new WorkerOffer("executor0", "host0", 4), - new WorkerOffer("executor1", "host1", 4)) + new WorkerOffer("executor0", "host0", 2), + new WorkerOffer("executor1", "host1", 1)) - // Stage 0 and stage 1 both in the DEFAULT profile; launch 2 tasks of each so both have running - // tasks. Distinct stageIds let us exclude one and count the other. - taskScheduler.submitTasks(FakeTask.createTaskSet(2, stageId = 0, stageAttemptId = 0)) - taskScheduler.submitTasks(FakeTask.createTaskSet(2, stageId = 1, stageAttemptId = 0)) + // Stage 0 and stage 1 both in the DEFAULT profile, 4 tasks each (8 total) but only 3 cores, so + // 3 launch and 5 remain enqueued. Distinct stageIds let us exclude one and count the other. + taskScheduler.submitTasks(FakeTask.createTaskSet(4, stageId = 0, stageAttemptId = 0)) + taskScheduler.submitTasks(FakeTask.createTaskSet(4, stageId = 1, stageAttemptId = 0)) val launched = taskScheduler.resourceOffers(workerOffers).flatten - assert(launched.length === 4, "all four tasks should launch (4 cores per executor)") + assert(launched.length === 3, "only three tasks can run (3 cores); the other five are enqueued") val defaultRp = ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID - // Nothing excluded: both stages' 4 tasks count. - assert(taskScheduler.runningTasksForOtherWorkInProfile(defaultRp, Set.empty) === 4) - // Exclude stage 0 (as if it were the group's own member): only stage 1's 2 tasks remain. - assert(taskScheduler.runningTasksForOtherWorkInProfile(defaultRp, Set(0)) === 2) + // Nothing excluded: all 8 tasks are outstanding (3 running + 5 enqueued) though only 3 run. + // The old running-only count would have reported 3 here; outstanding demand is what can hang a + // co-scheduled group, so it must be 8. + assert(taskScheduler.outstandingTasksForOtherWorkInProfile(defaultRp, Set.empty) === 8) + // Exclude stage 0 (as if it were the group's own member): only stage 1's 4 tasks remain. + assert(taskScheduler.outstandingTasksForOtherWorkInProfile(defaultRp, Set(0)) === 4) // Exclude both: zero "other" tasks. - assert(taskScheduler.runningTasksForOtherWorkInProfile(defaultRp, Set(0, 1)) === 0) - // A DIFFERENT (non-existent here) resource profile has no running tasks -- other-profile load - // must not be charged against this profile (finding A: RP-scoping). + assert(taskScheduler.outstandingTasksForOtherWorkInProfile(defaultRp, Set(0, 1)) === 0) + // A DIFFERENT (non-existent here) resource profile has no outstanding tasks -- other-profile + // load must not be charged against this profile (RP-scoping). val otherRpId = defaultRp + 12345 - assert(taskScheduler.runningTasksForOtherWorkInProfile(otherRpId, Set.empty) === 0, + assert(taskScheduler.outstandingTasksForOtherWorkInProfile(otherRpId, Set.empty) === 0, "tasks in the default profile must not be counted against a different profile") } From c136860a707c822bf73cebadf9ff50b75afc12b2 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Tue, 14 Jul 2026 04:28:39 +0000 Subject: [PATCH 05/54] [SPARK-XXXXX][CORE] Defer completion of pipelined-group members until the group finishes Group-observable completion for pipelined groups (spec S5). A stage co-scheduled with a still-running pipelined producer (a pipelined consumer) must not have its successful completion processed early: doing so would advance job completion and cancel the still-running producer (via cancelRunningIndependentStages), or make the consumer's output observable before the producer's. handleTaskCompletion: near the top, before any of the event's side effects, if a Success event belongs to a pipelined consumer with unfinished pipelined producers, buffer the whole CompletionEvent and return. This defers the entire event (coarse model), so its side effects (accumulator update, TaskEnd listener event, stage/job completion) run exactly once -- at replay. markStageAsFinished: when a pipelined producer finishes, release its deferred consumers -- but only when the producer's outcome is final. A ShuffleMapStage that finished without error yet is not isAvailable (an output missing; it will be resubmitted) is NOT treated as done, so its consumers stay deferred until the reattempt makes it available. On genuine success the buffered events are replayed; on producer failure they are dropped (the group reruns, S6) but their TaskEnd events are still emitted so listeners do not leak active-task accounting. cleanupStateForJobAndIndependentStages: drop any deferral keyed on a removed stage and remove it from other consumers' pending-producer sets, so no deferral outlives its job (e.g. on abort). assertDataStructuresEmpty also checks the deferral map is empty. Inert for jobs with no pipelined dependency: the deferral map is never populated, the completion check is a always-miss map lookup, and release/cleanup are no-ops. Tests (DAGSchedulerSuite): early-finishing consumer does not end the job or cancel its running producer; normal producer-then-consumer ordering; buffered completions dropped when the producer fails; TaskEnd fired exactly once (no buffer+replay duplication); deferral released only when the producer is genuinely available. Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 118 +++++++++- .../spark/scheduler/DAGSchedulerSuite.scala | 209 ++++++++++++++---- 2 files changed, 288 insertions(+), 39 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 89df4a5bad888..ee8b8ee271365 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -169,6 +169,24 @@ private[spark] class DAGScheduler( // Stages that must be resubmitted due to fetch failures private[scheduler] val failedStages = new HashSet[Stage] + /** + * Deferred completion for pipelined groups. When a stage is co-scheduled with a pipelined + * producer that is still running (a "pipelined consumer"), its successful task-completion events + * must not be processed yet: doing so could finish the consumer's job and cancel the + * still-running producer, or make the consumer's output observable before the producer's (spec + * S5, group-observable completion). We buffer such events here and replay them once every + * pipelined producer the consumer was waiting on has finished. + * + * Keyed by the consumer stage. `pendingProducers` is the set of its pipelined producer stages not + * yet finished; `bufferedEvents` are its Success CompletionEvents held until then. Entries exist + * only for stages that were co-scheduled with a not-yet-finished pipelined producer, so this is + * empty for any job without a pipelined dependency. + */ + private[scheduler] case class DeferredCompletion( + pendingProducers: HashSet[Stage] = new HashSet[Stage], + bufferedEvents: ListBuffer[CompletionEvent] = new ListBuffer[CompletionEvent]) + private[scheduler] val pipelinedConsumerDeferrals = new HashMap[Stage, DeferredCompletion] + private[scheduler] val activeJobs = new HashSet[ActiveJob] // Track all the jobs submitted by the same query execution, will clean up after @@ -1062,6 +1080,11 @@ private[spark] class DAGScheduler( logDebug("Removing stage %d from failed set.".format(stageId)) failedStages -= stage } + // Drop any pipelined-completion deferral keyed on this stage (as consumer), and + // remove it from other consumers' pending-producer sets (as producer), so no + // deferral outlives its job (e.g. on job abort before the producers finished). + pipelinedConsumerDeferrals -= stage + pipelinedConsumerDeferrals.values.foreach(_.pendingProducers -= stage) } // data structures based on StageId stageIdToStage -= stageId @@ -1921,6 +1944,11 @@ private[spark] class DAGScheduler( case None => 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, + // so its successful completions are deferred until those producers finish (S5). + val deferral = + pipelinedConsumerDeferrals.getOrElseUpdate(stage, DeferredCompletion()) + deferral.pendingProducers ++= pipelinedMissing submitMissingTasks(stage, jobId.get) // This stage is now running; if it is itself the pipelined producer of a waiting // consumer, co-schedule that consumer too. @@ -2719,6 +2747,29 @@ private[spark] class DAGScheduler( val stage = stageIdToStage(task.stageId) + // Group-observable completion (S5): if this stage is a pipelined consumer co-scheduled with a + // still-running pipelined producer, defer its *successful* completion in full until the + // producer(s) finish. Buffer the whole CompletionEvent and return before ANY of its side + // effects run (accumulator update, task-end listener event, stage/job completion) -- otherwise + // a consumer finishing ahead of its producer would advance job completion and cancel the + // still-running producer, or expose its output early. Deferring the entire event (not just the + // completion bookkeeping) is the coarse model, and crucially it makes the side effects run + // exactly ONCE, at replay: markStageAsFinished re-posts the buffered event once the last + // producer completes (or drops it if a producer fails, S6), and it then re-enters here and runs + // the side effects normally. This deferral check must therefore precede updateAccumulators and + // postTaskEnd. Inert for jobs with no pipelined dependency (the map is empty). + if (event.reason == Success) { + pipelinedConsumerDeferrals.get(stage) match { + case Some(deferral) if deferral.pendingProducers.nonEmpty => + logInfo(log"Deferring completion of task ${MDC(TASK_ID, event.taskInfo.taskId)} in " + + log"pipelined consumer ${MDC(STAGE, stage)} until its producer(s) " + + log"${MDC(MISSING_PARENT_STAGES, deferral.pendingProducers.toSeq)} finish") + deferral.bufferedEvents += event + return + case _ => + } + } + // Make sure the task's accumulators are updated before any other processing happens, so that // we can post a task end event before any jobs or stages are updated. The accumulators are // only updated in certain cases. @@ -3744,8 +3795,12 @@ private[spark] class DAGScheduler( /** * Marks a stage as finished and removes it from the list of running stages. + * + * `private[scheduler]` (not `private`) so tests can drive the pipelined-consumer release/retain + * decision below directly, including the retained branch (a not-yet-available pipelined producer + * about to resubmit), which is otherwise brittle to reach through the mock backend. */ - private def markStageAsFinished( + private[scheduler] def markStageAsFinished( stage: Stage, errorMessage: Option[String] = None, willRetry: Boolean = false): Unit = { @@ -3774,6 +3829,67 @@ private[spark] class DAGScheduler( } listenerBus.post(SparkListenerStageCompleted(stage.latestInfo)) runningStages -= stage + + // Release any pipelined consumers whose completion was deferred while this stage (a pipelined + // producer) was running. Only act when the producer's outcome is final: + // - willRetry: the stage is being retried, not finished -- leave consumers deferred. + // - a ShuffleMapStage that finished "successfully" (no errorMessage) but is NOT yet available + // (e.g. a bogus-epoch task left an output missing) is about to be resubmitted by + // processShuffleMapStageCompletion -- it is not truly done, so do NOT replay its consumers + // against soon-to-be-recomputed output; the release happens when its reattempt completes and + // it becomes available. + // Otherwise: producerFailed = the stage failed (errorMessage set) -> drop the consumers' + // buffered successes; producer succeeded -> replay them. + if (!willRetry) { + val producerFailed = errorMessage.isDefined + val producerAboutToResubmit = stage match { + case m: ShuffleMapStage => !producerFailed && !m.isAvailable + case _ => false + } + if (!producerAboutToResubmit) { + releaseDeferredPipelinedConsumers(stage, producerFailed = producerFailed) + } + } + } + + /** + * Called when a stage finishes. If `finishedStage` is a pipelined producer that some co-scheduled + * consumer was deferred on, remove it from that consumer's pending-producer set. When a consumer + * has no pending producers left, either replay its buffered completion events (producer succeeded) + * or drop them (a producer failed -- the group will be torn down and rerun, so the consumer's + * buffered successes must not be applied; S6). Inert unless `finishedStage` is a tracked producer. + */ + private def releaseDeferredPipelinedConsumers( + finishedStage: Stage, producerFailed: Boolean): Unit = { + if (pipelinedConsumerDeferrals.isEmpty) { + return + } + // Consumers waiting on this producer. + val released = pipelinedConsumerDeferrals.filter { + case (_, d) => d.pendingProducers.contains(finishedStage) + }.keys.toArray + for (consumer <- released) { + val deferral = pipelinedConsumerDeferrals(consumer) + deferral.pendingProducers -= finishedStage + if (deferral.pendingProducers.isEmpty) { + pipelinedConsumerDeferrals -= consumer + val events = deferral.bufferedEvents.toList + if (producerFailed) { + logInfo(log"Dropping ${MDC(NUM_EVENTS, events.size.toLong)} deferred completion(s) for " + + log"pipelined consumer ${MDC(STAGE, consumer)} because its producer " + + log"${MDC(STAGE, finishedStage)} failed; the group will be rerun") + // The buffered tasks genuinely succeeded, so still emit their TaskEnd events -- otherwise + // listeners that track active tasks (e.g. AppStatusListener) would believe these tasks + // are still running. We deliberately do NOT run the stage/job completion bookkeeping: + // the group failed and will be rerun, so the consumer's results must not be applied. + events.foreach(postTaskEnd) + } else { + logInfo(log"Replaying ${MDC(NUM_EVENTS, events.size.toLong)} deferred completion(s) for " + + log"pipelined consumer ${MDC(STAGE, consumer)} now that its producers have finished") + events.foreach(eventProcessLoop.post) + } + } + } } /** 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 2bf567d4c8c9d..2e00407aca929 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6379,44 +6379,14 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } - test("pipelined shuffle: one producer feeding two pipelined consumers reconsiders both") { - // R --regular--> P --pipelined--> C1 and C2 ; C1,C2 --regular--> D. - // When P starts (R done), reconsideration must resubmit BOTH parked consumers C1 and C2. - // (Guards the `.filter` in submitWaitingPipelinedChildStages against a `.find` regression.) - val rddR = new MyRDD(sc, 2, Nil) - val rddP = new MyRDD(sc, 2, List(new ShuffleDependency(rddR, new HashPartitioner(2))), - tracker = mapOutputTracker) - val rddC1 = new MyRDD(sc, 2, List(new PipelinedShuffleDependency(rddP, new HashPartitioner(2))), - tracker = mapOutputTracker) - val rddC2 = new MyRDD(sc, 2, List(new PipelinedShuffleDependency(rddP, new HashPartitioner(2))), - tracker = mapOutputTracker) - val rddD = new MyRDD(sc, 2, - List(new ShuffleDependency(rddC1, new HashPartitioner(2)), - new ShuffleDependency(rddC2, new HashPartitioner(2))), - tracker = mapOutputTracker) - submit(rddD, Array(0, 1)) - - // Only R runs initially; P, C1, C2 wait (D waits on its regular parents C1, C2). - assert(taskSets.size === 1, s"expected only R submitted, got ${taskSets.size}") - - // R completes -> P runs and reconsiders BOTH C1 and C2. - val idR = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddR).get.stageId - completeShuffleMapStageSuccessfully(idR, 0, 2) - assert(scheduler.runningStages.exists(_.rdd eq rddC1), "C1 should be co-scheduled with P") - assert(scheduler.runningStages.exists(_.rdd eq rddC2), "C2 should be co-scheduled with P") - - // Drain P, then C1, C2, then D. - val idP = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddP).get.stageId - completeShuffleMapStageSuccessfully(idP, 0, 2) - val idC1 = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddC1).get.stageId - completeShuffleMapStageSuccessfully(idC1, 0, 2) - val idC2 = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddC2).get.stageId - completeShuffleMapStageSuccessfully(idC2, 0, 2) - val tsD = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddD).get - complete(tsD, Seq((Success, 42), (Success, 43))) - assert(results === Map(0 -> 42, 1 -> 43)) - assertDataStructuresEmpty() - } + // Note: there is deliberately no "one producer feeding two pipelined consumers" test. That shape + // (a single producer stage with two pipelined consumers) is pipelined *fan-out*, which requires a + // shared PipelinedShuffleDependency and is forbidden by the spec (S9); its admission-time + // rejection is a later milestone. Two distinct PipelinedShuffleDependency objects over one RDD do + // NOT create it -- they create two separate producer stages, each with one consumer. So for every + // supported (non-fan-out) shape a running pipelined producer has at most one waiting pipelined + // consumer, and submitWaitingPipelinedChildStages reconsiders it via the single-child path + // covered by the transitive-cascade and producer-behind-regular-shuffle tests. test("pipelined shuffle: reconsideration cascades transitively (B->C->D behind a regular root)") { // Z --regular--> B --pipelined--> C --pipelined--> D. When Z completes, B runs and reconsiders @@ -6728,6 +6698,169 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti myScheduler.outstandingTasksForOtherWorkForTest = (_, _) => 0 } } + + // ========================================================================================== + // Group-observable (coarse deferred) completion (spec S5) + // ========================================================================================== + + test("pipelined shuffle: a consumer that finishes early does not end the job or cancel its " + + "still-running producer") { + // producer (shuffle map) --[pipelined]--> consumer (result). Consumer is co-scheduled; if its + // result tasks finish before the producer, its completion MUST be deferred -- otherwise the job + // would end and cancelRunningIndependentStages would cancel the still-running 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) + submit(consumerRdd, Array(0, 1)) + assert(taskSets.size === 2) + val producerStageId = taskSets.head.stageId + val consumerTaskSet = taskSets(1) + + // Consumer finishes FIRST (both result tasks succeed) while the producer is still running. + complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) + // The job must NOT be finished yet, and the producer must still be running (not cancelled). + assert(results.isEmpty, "consumer completion must be deferred until the producer finishes") + assert(scheduler.runningStages.exists(_.rdd eq producerRdd), + "the still-running producer must not be cancelled by the early-finishing consumer") + assert(cancelledStages.isEmpty) + + // Now the producer finishes -> the deferred consumer completions replay -> the job completes. + completeShuffleMapStageSuccessfully(producerStageId, 0, 2) + assert(results === Map(0 -> 42, 1 -> 43), + "deferred consumer completions should replay once the producer finishes") + assertDataStructuresEmpty() + } + + test("pipelined shuffle: producer-then-consumer ordering completes normally (no deferral " + + "needed)") { + // Sanity: when the producer finishes before the consumer's tasks complete, there is nothing to + // defer and the job completes as usual. + 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 consumerTaskSet = taskSets(1) + + completeShuffleMapStageSuccessfully(producerStageId, 0, 2) + complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: deferred consumer completions are dropped when the producer fails") { + // If a producer fails while a co-scheduled consumer's completions are buffered, those buffered + // successes must be DROPPED (not applied): the group is torn down / the job fails, and applying + // a partial consumer success would be incorrect (S6). + 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 producerTaskSet = taskSets.head + val consumerTaskSet = taskSets(1) + + // Consumer finishes early -> its completions are buffered (deferred). + complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) + assert(results.isEmpty, "consumer completions should be buffered while the producer runs") + + // The producer now FAILS its whole task set. The job fails; the buffered consumer successes + // must NOT have been applied as results. + failed(producerTaskSet, "producer blew up") + assert(failure.get() != null, "job should fail when the producer fails") + assert(results.isEmpty, + "buffered consumer successes must be dropped when the producer fails, not applied") + 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 + // update) applied exactly once -- at replay -- not once when buffered and again when replayed. + // The producer (2 tasks) and consumer (2 tasks) yield exactly 4 TaskEnd events total; a + // buffer+replay double-post would inflate that count. + 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) + submit(consumerRdd, Array(0, 1)) + val producerStageId = taskSets.head.stageId + val consumerTaskSet = taskSets(1) + + // Consumer finishes first -> its two completions are buffered (deferred, no TaskEnd yet). + complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) + // Producer finishes -> the two deferred consumer completions replay. + completeShuffleMapStageSuccessfully(producerStageId, 0, 2) + assert(results === Map(0 -> 42, 1 -> 43)) + sc.listenerBus.waitUntilEmpty(10000) + + // Exactly 4 TaskEnd events (2 producer + 2 consumer). A double-post from buffer+replay would + // make it 6 (the 2 consumer tasks counted twice). + assert(taskEndCount.get() === 4, + s"expected 4 TaskEnd events (no buffer+replay duplication), got ${taskEndCount.get()}") + assertDataStructuresEmpty() + } finally { + sc.removeSparkListener(countingListener) + } + } + + test("pipelined shuffle: releasing a deferral is gated on the producer being genuinely done") { + // Regression for the "producer finished but not available -> resubmitted" hazard: a + // ShuffleMapStage producer that reaches markStageAsFinished with no error but is NOT available + // (some output missing, so it will be resubmitted) must NOT release/replay its deferred + // consumers -- doing so would apply the consumer's results against soon-to-be-recomputed + // output. + // This drives BOTH branches of the gate: (1) the RETAINED branch -- markStageAsFinished on a + // not-yet-available pipelined producer must leave the deferral in place; and (2) the RELEASED + // branch -- once the producer is genuinely available, the deferral is released and results + // applied. + 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 consumerTaskSet = taskSets(1) + + complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) + // Deferred while the producer runs. + assert(scheduler.pipelinedConsumerDeferrals.keys.exists(_.rdd eq consumerRdd)) + assert(results.isEmpty) + + val producerStage = + scheduler.stageIdToStage(producerStageId).asInstanceOf[ShuffleMapStage] + + // (1) RETAINED branch: the producer has completed NO partitions yet, so it is not available. + // markStageAsFinished with no error and willRetry=false hits the producerAboutToResubmit guard + // (ShuffleMapStage && !producerFailed && !isAvailable) and must NOT release the consumer. + assert(!producerStage.isAvailable, "precondition: producer not yet available") + scheduler.markStageAsFinished(producerStage, errorMessage = None, willRetry = false) + assert(scheduler.pipelinedConsumerDeferrals.keys.exists(_.rdd eq consumerRdd), + "a not-yet-available pipelined producer must NOT release its deferred consumers (it is " + + "about to resubmit; releasing would apply results against soon-to-be-recomputed output)") + assert(results.isEmpty, + "consumer results must not be applied while the producer is not available") + + // (2) RELEASED branch: producer completes for real and IS available -> deferral released, + // consumer results applied. (Re-add it to runningStages since (1) removed it, so the normal + // completion path proceeds as in production.) + scheduler.runningStages += producerStage + completeShuffleMapStageSuccessfully(producerStageId, 0, 2) + assert(producerStage.isAvailable) + assert(!scheduler.pipelinedConsumerDeferrals.keys.exists(_.rdd eq consumerRdd), + "deferral released once the producer is genuinely done (available)") + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + } class DAGSchedulerAbortStageOffSuite extends DAGSchedulerSuite { From 768c398565bb0e25348a84390a0588ff4f12b644 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sat, 18 Jul 2026 02:41:28 +0000 Subject: [PATCH 06/54] [SPARK-XXXXX][CORE] Fix pipelined-group admission: skip zombie attempts, reject dynamic allocation Two correctness fixes to the pipelined-group slot admission, found in review: - Skip zombie attempts in outstandingTasksForOtherWorkInProfile. The occupancy count summed numTasks - tasksSuccessful over every attempt in taskSetsByStageIdAndAttempt, including zombie (superseded) attempts. A retried/killed stage can have both a zombie and a live attempt in the map at once, and the live attempt already re-runs the zombie's outstanding tasks -- so counting both double-counted that stage's demand, inflating occupancy and potentially failing a pipelined group with CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT that would actually fit. Add the !isZombie filter used elsewhere on this map. - Reject a pipelined-shuffle job under dynamic allocation. A pipelined group is gang-scheduled and its free-slot check measures currently-active executors; under dynamic allocation a job can be submitted before any executor has spun up, so the group would be failed against a transient 0-slot snapshot even though the cluster would soon have capacity. Barrier scheduling forbids the same combination (checkBarrierStageWithDynamicAllocation); pipelined groups now do likewise. The former rejectSpeculationWithPipelinedShuffle is generalized to rejectUnsupportedPipelinedJob, which rejects both speculation and dynamic allocation (single RDD-graph walk, run only when a relevant feature is enabled). Tests: a zombie + live attempt is counted once, not twice; a pipelined job under dynamic allocation is rejected while a regular job under dynamic allocation is not. Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 67 +++++++++++++------ .../spark/scheduler/TaskSchedulerImpl.scala | 13 ++-- .../spark/scheduler/DAGSchedulerSuite.scala | 45 +++++++++++++ .../scheduler/TaskSchedulerImplSuite.scala | 24 +++++++ 4 files changed, 123 insertions(+), 26 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 ee8b8ee271365..9d65669cece0c 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -985,28 +985,50 @@ private[spark] class DAGScheduler( } /** - * Reject a job that enables speculation and uses a pipelined shuffle: a speculative copy of a - * producer would race a consumer already reading the producer's partial output, with no commit - * barrier protecting the read (spec S9). Checked up front, before any stage is created, so a - * rejection leaves no partial scheduler state. Used by the result-job path (handleJobSubmitted); - * the map-stage-job path rejects a pipelined dependency outright (see handleMapStageSubmitted), - * which subsumes speculation. Returns true (and fails the job via `listener`) if rejected; false - * otherwise. Inert for jobs without a pipelined dependency. + * Reject a job that uses a pipelined shuffle in combination with a cluster feature that a + * pipelined group cannot support. Checked up front, before any stage is created, so a rejection + * leaves no partial scheduler state. Used by the result-job path (handleJobSubmitted); the + * map-stage-job path rejects a pipelined dependency outright (see handleMapStageSubmitted), which + * subsumes these. Returns true (and fails the job via `listener`) if rejected; false otherwise. + * The RDD-graph walk runs only when a relevant feature is enabled and is inert for jobs without a + * pipelined dependency. + * + * Rejected combinations: + * - Speculation: a speculative copy of a producer would race a consumer already reading the + * producer's partial output, with no commit barrier protecting the read (spec S9). + * - Dynamic allocation: a pipelined group is gang-scheduled (all member stages must run at + * once), and the free-slot admission check measures currently-active executors. Under dynamic + * allocation a job can start before any executor has spun up, so the group would be failed + * with CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT against a transient 0-slot snapshot even though + * the cluster would soon have capacity. Barrier scheduling forbids the same combination + * (checkBarrierStageWithDynamicAllocation); pipelined groups do likewise. */ - private def rejectSpeculationWithPipelinedShuffle( + private def rejectUnsupportedPipelinedJob( jobId: Int, finalRDD: RDD[_], listener: JobListener): Boolean = { - if (sc.conf.get(config.SPECULATION_ENABLED) && rddGraphHasPipelinedDependency(finalRDD)) { - logWarning(log"Rejecting job ${MDC(JOB_ID, jobId)}: speculation is incompatible with " + - log"pipelined shuffle dependencies") + // Only walk the RDD graph if a relevant feature is on, then only once. + val speculationOn = sc.conf.get(config.SPECULATION_ENABLED) + val dynAllocOn = Utils.isDynamicAllocationEnabled(sc.conf) + if ((speculationOn || dynAllocOn) && rddGraphHasPipelinedDependency(finalRDD)) { + if (speculationOn) { + logWarning(log"Rejecting job ${MDC(JOB_ID, jobId)}: speculation is incompatible with " + + log"pipelined shuffle dependencies") + listener.jobFailed(new SparkException( + "Speculative execution is not supported for a job that uses a pipelined shuffle " + + s"dependency. Disable ${config.SPECULATION_ENABLED.key} for such jobs.")) + return true + } + // dynAllocOn + logWarning(log"Rejecting job ${MDC(JOB_ID, jobId)}: dynamic allocation is incompatible " + + log"with pipelined shuffle dependencies") listener.jobFailed(new SparkException( - "Speculative execution is not supported for a job that uses a pipelined shuffle " + - s"dependency. Disable ${config.SPECULATION_ENABLED.key} for such jobs.")) - true - } else { - false + "Dynamic allocation is not supported for a job that uses a pipelined shuffle dependency: " + + "a pipelined stage group must be co-scheduled against a known cluster size. Disable " + + s"${config.DYN_ALLOCATION_ENABLED.key} for such jobs.")) + return true } + false } /** Invoke `.partitions` on the given RDD and all of its ancestors */ @@ -1724,12 +1746,13 @@ private[spark] class DAGScheduler( return } - // A job that uses a pipelined shuffle runs its producer and consumer stages concurrently, so a - // speculative copy of a producer would race a consumer already reading the producer's partial - // output, with no commit barrier protecting the read. Reject such a job up front -- before any - // stages are created, so no partial scheduler state is left behind. (Inert for jobs without - // any pipelined dependency.) - if (rejectSpeculationWithPipelinedShuffle(jobId, finalRDD, listener)) { + // A job that uses a pipelined shuffle co-schedules its producer and consumer stages, which is + // incompatible with speculation (a speculative producer copy would race a consumer reading its + // partial output) and with dynamic allocation (the gang-scheduling slot check would fail + // against a not-yet-warmed cluster). Reject such a job up front -- before any stages are + // created, so no partial scheduler state is left behind. (Inert for jobs without any pipelined + // dependency.) + if (rejectUnsupportedPipelinedJob(jobId, finalRDD, listener)) { return } diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala index fba876717cef8..e162a5c19cb69 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala @@ -169,9 +169,10 @@ private[spark] class TaskSchedulerImpl( * by everything else in the SAME resource profile, so it excludes the group's own members (whose * tasks would otherwise be charged against the group's own admission). * - * Computed from the TaskSetManagers so it is resource-profile-scoped, and taken under a single - * lock so the count is one consistent snapshot (both the per-profile total and the excluded - * members are read together). + * Computed from the TaskSetManagers so it is resource-profile-scoped, skips zombie (superseded) + * attempts so a retried stage is not double-counted, and is taken under a single lock so the + * count is one consistent snapshot (both the per-profile total and the excluded members are read + * together). */ def outstandingTasksForOtherWorkInProfile( resourceProfileId: Int, excludeStageIds: Set[Int]): Int = synchronized { @@ -180,7 +181,11 @@ private[spark] class TaskSchedulerImpl( Iterator.empty } else { attempts.valuesIterator - .filter(_.taskSet.resourceProfileId == resourceProfileId) + // Skip zombie attempts (superseded by a retry/kill): a stage can have both a zombie and a + // live attempt in this map at once, and the live attempt already re-runs the zombie's + // outstanding tasks -- counting both would double-count that stage's demand. Matches the + // !isZombie filtering used elsewhere on this map. + .filter(tsm => !tsm.isZombie && tsm.taskSet.resourceProfileId == resourceProfileId) .map(tsm => math.max(0, tsm.numTasks - tsm.tasksSuccessful)) } }.sum 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 2e00407aca929..32971ca3a140c 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6492,6 +6492,51 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } } + test("pipelined shuffle: dynamic allocation is rejected for a job with a pipelined dependency") { + // A pipelined group is gang-scheduled and its slot check measures currently-active executors; + // under dynamic allocation a job can start before any executor is up, so the check would fail + // the group against a transient 0-slot snapshot. Reject the combo (barrier does the same). + // DYN_ALLOCATION_TESTING lets isDynamicAllocationEnabled report true under a local master. + sc.conf.set(config.DYN_ALLOCATION_ENABLED, true) + sc.conf.set(config.DYN_ALLOCATION_TESTING, true) + 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 = {} + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(consumerRdd, Array(0, 1), listener = failListener) + assert(failure.get() != null, + "job with pipelined dependency + dynamic allocation should be rejected") + assert(failure.get().getMessage.contains("Dynamic allocation is not supported")) + assert(taskSets.isEmpty) + assertDataStructuresEmpty() + } finally { + sc.conf.set(config.DYN_ALLOCATION_ENABLED, false) + sc.conf.set(config.DYN_ALLOCATION_TESTING, false) + } + } + + 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. + sc.conf.set(config.DYN_ALLOCATION_ENABLED, true) + sc.conf.set(config.DYN_ALLOCATION_TESTING, true) + try { + val shuffleMapRdd = new MyRDD(sc, 2, Nil) + val shuffleDep = new ShuffleDependency(shuffleMapRdd, new HashPartitioner(2)) + val reduceRdd = new MyRDD(sc, 2, List(shuffleDep), tracker = mapOutputTracker) + submit(reduceRdd, Array(0, 1)) + assert(taskSets.nonEmpty, "regular-shuffle job must not be rejected under dynamic allocation") + } finally { + sc.conf.set(config.DYN_ALLOCATION_ENABLED, false) + sc.conf.set(config.DYN_ALLOCATION_TESTING, false) + } + } + test("pipelined shuffle: submitting a pipelined dependency as a map-stage job is rejected") { // A map-stage job (SparkContext.submitMapStage, e.g. AQE's ShuffleExchangeExec) materializes a // shuffle to produce map-output statistics. A PipelinedShuffleDependency cannot serve that -- diff --git a/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala index ef21af6096487..edf2dcdc0e75a 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala @@ -2786,4 +2786,28 @@ class TaskSchedulerImplSuite extends SparkFunSuite with LocalSparkContext "tasks in the default profile must not be counted against a different profile") } + test("outstandingTasksForOtherWorkInProfile does not double-count a zombie + live attempt") { + // A retried/killed stage can have both a zombie (superseded) attempt and a live attempt in the + // map at once; the live attempt re-runs the zombie's outstanding tasks, so only the live + // attempt's demand should count. Counting both would inflate occupancy and could spuriously + // fail a pipelined group with CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT. + val taskScheduler = setupScheduler() + val defaultRp = ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID + + // Attempt 0 of stage 0: 4 tasks, none started. Mark it zombie (as if superseded). + val attempt0 = FakeTask.createTaskSet(4, stageId = 0, stageAttemptId = 0) + taskScheduler.submitTasks(attempt0) + taskScheduler.taskSetManagerForAttempt(0, 0).get.isZombie = true + + // Attempt 1 of the SAME stage: the live retry, also 4 tasks. + val attempt1 = FakeTask.createTaskSet(4, stageId = 0, stageAttemptId = 1) + taskScheduler.submitTasks(attempt1) + + // Both attempts are in the map, but only the live attempt's 4 tasks are outstanding -- not 8. + assert(taskScheduler.outstandingTasksForOtherWorkInProfile(defaultRp, Set.empty) === 4, + "a zombie attempt must not be counted alongside its live retry") + // Excluding the stage (as the group's own member) drops both attempts. + assert(taskScheduler.outstandingTasksForOtherWorkInProfile(defaultRp, Set(0)) === 0) + } + } From d3749e244bc4737f13f8f5d92fd7d293b9e1d059 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sat, 18 Jul 2026 06:20:32 +0000 Subject: [PATCH 07/54] [SPARK-XXXXX][CORE] Document that isPipelinedGroupMember is consumed by a later PR in the stack isPipelinedGroupMember is defined here alongside the other group-topology helpers but is first used by the group-atomic failure handling added in a later PR of this stack (TaskSet.isPipelined and the member-failure fail-fast, spec S6). Reviewed in isolation this change reads as introducing an unused private method; note the forward reference in the scaladoc so it is not mistaken for dead code. Co-authored-by: Isaac --- .../main/scala/org/apache/spark/scheduler/DAGScheduler.scala | 5 +++++ 1 file changed, 5 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 9d65669cece0c..488615d06e5fb 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -1583,6 +1583,11 @@ private[spark] class DAGScheduler( * failure must fail the whole group rather than resubmit one stage; the task scheduler keys its * fail-fast behavior off this (via TaskSet.isPipelined). False for any stage in a job with no * pipelined dependency. + * + * NOTE: this is deliberately defined here alongside the other group-topology helpers + * (`isPipelinedProducer`, `pipelinedGroupOf`) but is first consumed by the group-atomic failure + * handling added in a later PR of this stack (`TaskSet.isPipelined` and the member-failure + * fail-fast, spec S6). It reads as unused when this change is viewed alone; that is expected. */ private def isPipelinedGroupMember(stage: Stage): Boolean = { if (isPipelinedProducer(stage)) { From b63f9fe8b815827c4f4336935a40fb2f4f7e8b0f Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sat, 18 Jul 2026 17:33:35 +0000 Subject: [PATCH 08/54] [SPARK-XXXXX][CORE] Fix stale free-slots doc wording: outstanding = running plus enqueued The admission check counts OTHER work's OUTSTANDING tasks (running plus enqueued, `numTasks - tasksSuccessful`) against a resource profile's capacity -- as `outstandingTasksForOtherWork` / `TaskSchedulerImpl. outstandingTasksForOtherWorkInProfile` and the config doc already state. Two spots in DAGScheduler still described this as only the tasks "already running", which was stale after the count was widened to include enqueued tasks; the `pipelinedGroupExceedsCapacity` scaladoc even contradicted its own later "running or enqueued" line. Reword both to "outstanding -- running plus enqueued" so the prose matches the code. Doc/comment only; no behavior change. Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 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 488615d06e5fb..5e23580d42b16 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -1511,12 +1511,12 @@ private[spark] class DAGScheduler( * reservation, would deadlock -- so we fail fast with a clear error. * * Demand is compared against FREE slots, not total capacity (spec S4.1): free = total capacity - * minus the tasks already running for OTHER work (other groups and regular jobs) in the SAME - * resource profile. Comparing against total capacity would admit a group that fits the cluster in - * principle but cannot co-fit right now beside a busy neighbor, then hang. Occupancy excludes the - * group's OWN already-running members (e.g. a producer submitted just before its consumer is - * admitted) -- otherwise the group would be charged for its own slots and could fail a check it - * actually fits. + * minus what OTHER work (other groups and regular jobs) has outstanding -- running plus enqueued + * -- in the SAME resource profile. Comparing against total capacity would admit a group that + * fits the cluster in principle but cannot co-fit right now beside a busy neighbor, then hang. + * Occupancy excludes the group's OWN already-running members (e.g. a producer submitted just + * before its consumer is admitted) -- otherwise the group would be charged for its own slots + * and could fail a check it actually fits. * * Occupancy is resource-profile-scoped to match `totalSlots` (which is per-profile): counting * other profiles' running tasks against one profile's capacity would spuriously reject a fitting @@ -1957,9 +1957,9 @@ private[spark] class DAGScheduler( if (regularMissing.isEmpty && allPipelinedParentsRunning) { // Best-effort slot check: the whole pipelined group must run at once, so its demand - // must fit in the currently-FREE slots (total capacity minus what other work is - // already running; spec S4.1). If it cannot, fail fast with a clear error rather than - // deadlock (there is no out-of-band slot reservation in v1). + // must fit in the currently-FREE slots (total capacity minus what other work has + // outstanding -- running plus enqueued; spec S4.1). If it cannot, fail fast with a + // clear error rather than deadlock (there is no out-of-band slot reservation in v1). pipelinedGroupExceedsCapacity(stage) match { case Some((demand, freeSlots)) => abortStage(stage, s"Cannot co-schedule pipelined stage group: needs $demand " + From f30bf130e0ed5f264951be4da16411d70e971c3e Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sun, 19 Jul 2026 06:05:39 +0000 Subject: [PATCH 09/54] [SPARK-XXXXX][CORE] Restrict pipelined jobs to all-regular-or-all-PG; admit the group up front v1 (M1, real-time-mode scope) supports a job that is either all-regular or all-pipelined, not a mix. This lets an all-pipelined job's whole stage graph be one pipelined group with no regular prefix, so gang admission can be decided up front -- before any member stage is submitted -- rather than mid-DAG once a producer is already running. Changes in handleJobSubmitted (before createResultStage, so a rejection leaves no partial scheduler state, exactly like the barrier slot check and the speculation/DA reject): - classifyJobShuffleKinds walks the RDD graph and rejects a job that mixes a pipelined shuffle with a regular one, fail-fast. - rejectUnadmittablePipelinedGroup computes the whole all-pipelined group's concurrent-task demand from the RDD graph and checks it against free slots (maxNumConcurrentTasks minus other work's running-plus-enqueued demand, spec S4.1); fails the job if it cannot fit. No scheduler retry -- a transient shortfall is the caller's to retry (e.g. the streaming batch loop reruns the batch). This is true all-or-nothing gang admission: the whole group is admitted, or the job fails before any member runs, so a member is never left running while a sibling cannot get slots. The late slot check in submitStage's co-schedule branch (which measured a mid-flight snapshot after a producer was already running) is removed, along with the now-unused pipelinedGroupExceedsCapacity / pipelinedGroupOf; the capacity/occupancy seams are refactored to be resource-profile-keyed. Tests: mixed jobs rejected up front; all-PG group admitted/rejected up front (2- and 3-stage); the removed mid-DAG prefix tests are dropped as that shape is no longer supported. Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 248 ++++++++++------- .../spark/scheduler/DAGSchedulerSuite.scala | 261 ++++-------------- 2 files changed, 190 insertions(+), 319 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 5e23580d42b16..b83dc8ee555a7 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -984,6 +984,35 @@ private[spark] class DAGScheduler( } } + /** + * Classifies the shuffle boundaries in the RDD graph rooted at `finalRDD` by kind, walking the + * RDD dependency graph directly (before any stages are created). Returns whether the graph + * contains any [[PipelinedShuffleDependency]] and whether it contains any regular (non-pipelined) + * `ShuffleDependency`. Narrow dependencies are not boundaries and are ignored. + * + * v1 (M1, real-time-mode scope) supports a job that is either ALL-regular or ALL-pipelined: a + * job whose shuffle graph mixes a pipelined shuffle with a regular one is rejected fail-fast + * (see `handleJobSubmitted`). Restricting to a single kind makes the whole job one pipelined + * group when pipelined, so gang admission can be decided up front before any stage is submitted. + */ + private def classifyJobShuffleKinds(finalRDD: RDD[_]): (Boolean, Boolean) = { + var hasPipelined = false + var hasRegular = false + traverseRDDGraph(finalRDD) { (rdd, enqueue) => + rdd.dependencies.foreach { dep => + dep match { + case _: PipelinedShuffleDependency[_, _, _] => hasPipelined = true + case _: ShuffleDependency[_, _, _] => hasRegular = true + case _ => // narrow dependency: not a boundary + } + // Descend through every edge (shuffle and narrow) so a pipelined boundary behind a regular + // one -- or vice versa -- anywhere in the graph is still detected. traverseRDDGraph dedups. + enqueue(dep.rdd) + } + } + (hasPipelined, hasRegular) + } + /** * Reject a job that uses a pipelined shuffle in combination with a cluster feature that a * pipelined group cannot support. Checked up front, before any stage is created, so a rejection @@ -1472,106 +1501,97 @@ private[spark] class DAGScheduler( } /** - * The full pipelined group `stage` belongs to: the connected component of the stage graph over - * pipelined edges (both the producers `stage` reads through a [[PipelinedShuffleDependency]] and, - * transitively, their pipelined producers/consumers). Walks `parents` and the shuffle-map stages - * of pipelined dependencies. Returns just `stage` if it has no pipelined edge. + * The total concurrent-task demand of an all-pipelined job, computed from the RDD graph BEFORE + * any stage is created (so a rejection based on it leaves no partial scheduler state, exactly as + * the barrier slot check and the speculation/DA reject do). Because a v1 job is either + * all-regular or all-pipelined, an all-pipelined job's whole stage graph is one pipelined group; + * its members are the final result stage plus every pipelined producer. Each member's task count + * is its RDD's partition count (`rdd.partitions.length`), matching how `createShuffleMapStage` + * derives `numTasks`. `finalNumPartitions` is the result stage's task count (the number of + * partitions the job runs, which may be a subset of `finalRDD.partitions`). */ - private def pipelinedGroupOf(stage: Stage): Set[Stage] = { - val group = new HashSet[Stage] - val toVisit = new ListBuffer[Stage] - toVisit += stage - while (toVisit.nonEmpty) { - val s = toVisit.remove(0) - if (group.add(s)) { - // Pipelined producers this stage reads (direct pipelined parent shuffle-map stages). - s.parents.foreach { - case m: ShuffleMapStage - if m.shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]] => toVisit += m - case _ => - } - // Pipelined consumers/producers reachable the other direction: any known stage co-scheduled - // with s via a pipelined edge. We only have parent links, so also pull in stages whose - // pipelined parent is s (search the created stages). - stageIdToStage.valuesIterator.foreach { other => - if (!group.contains(other) && other.parents.contains(s) && - isPipelinedProducer(s)) { - toVisit += other - } - } + private def pipelinedJobConcurrentTaskDemand(finalRDD: RDD[_], finalNumPartitions: Int): Int = { + var demand = finalNumPartitions + traverseRDDGraph(finalRDD) { (rdd, enqueue) => + rdd.dependencies.foreach { + case pd: PipelinedShuffleDependency[_, _, _] => demand += pd.rdd.partitions.length + case _ => // regular/narrow deps do not occur in an all-pipelined job's group } + rdd.dependencies.foreach(dep => enqueue(dep.rdd)) } - group.toSet + demand } /** - * Best-effort admission check for a pipelined group: all its member stages must run concurrently, - * so the cluster must have enough task slots for the whole group at once. If the group's demand - * exceeds what is currently available it can never be co-resident and, lacking an out-of-band slot - * reservation, would deadlock -- so we fail fast with a clear error. - * - * Demand is compared against FREE slots, not total capacity (spec S4.1): free = total capacity - * minus what OTHER work (other groups and regular jobs) has outstanding -- running plus enqueued - * -- in the SAME resource profile. Comparing against total capacity would admit a group that - * fits the cluster in principle but cannot co-fit right now beside a busy neighbor, then hang. - * Occupancy excludes the group's OWN already-running members (e.g. a producer submitted just - * before its consumer is admitted) -- otherwise the group would be charged for its own slots - * and could fail a check it actually fits. + * Up-front gang admission for an all-pipelined job (v1 / M1), checked BEFORE any stage exists. + * Because a v1 job is either all-regular or all-pipelined (mixed jobs are already rejected), an + * all-pipelined job's whole stage graph is one pipelined group with no regular prefix, so its + * full demand is known up front and the group is ready to admit immediately. Checking here + * (rather than in `submitStage` once a producer is already running) is true all-or-nothing gang + * admission: the whole group is admitted, or the job is failed before any member runs, so a + * member is never left running while a sibling cannot get slots -- and, like the barrier slot + * check, a rejection leaves no partial scheduler state. * - * Occupancy is resource-profile-scoped to match `totalSlots` (which is per-profile): counting - * other profiles' running tasks against one profile's capacity would spuriously reject a fitting - * group whenever an unrelated profile is busy. v1 requires a group to be single-profile (spec S9), - * so the whole group shares `stage`'s profile. - * - * This is best-effort, checked once when the group is first co-scheduled; it is NOT the atomic - * gang reservation deferred to a later hardening step. Without that reservation a race between two - * groups admitting concurrently can still transiently over-admit (each sees the other's slots as - * free before either's tasks register as running); the reservation closes that race later. What - * this does guarantee now is that a group is never admitted against slots a busy neighbor is - * already committed to (running or enqueued), which is the common single-group-on-a-busy-cluster - * hang. + * Free-slot accounting follows spec S4.1: demand vs. total capacity (`maxNumConcurrentTasks` for + * the default profile -- v1 requires a single-profile group, spec S9) minus what OTHER work + * (other jobs) has outstanding -- running plus enqueued -- in that profile. Counting enqueued, + * not just running, tasks charges a busy neighbor's queued backlog against capacity too, so a + * group is not admitted against slots other work is already committed to. Fails the job via + * `listener` with no retry -- a transient shortfall is the caller's to retry (e.g. the streaming + * batch loop reruns the batch). Returns true (job failed) if it does not fit. * * The check can be turned off with `spark.scheduler.pipelinedGroup.slotCheck.enabled=false` (for - * deployments that admit capacity out-of-band, e.g. via a slot reservation), in which case this - * always returns None and the group is co-scheduled unconditionally. - * - * Returns Some((demand, freeSlots)) when the group does not fit, else None. + * deployments that admit capacity out-of-band, e.g. via a slot reservation). */ - private def pipelinedGroupExceedsCapacity(stage: Stage): Option[(Int, Int)] = { + private def rejectUnadmittablePipelinedGroup( + jobId: Int, finalRDD: RDD[_], partitions: Array[Int], listener: JobListener): Boolean = { if (!sc.conf.get(config.PIPELINED_GROUP_SLOT_CHECK_ENABLED)) { - return None - } - val group = pipelinedGroupOf(stage) - val demand = group.toSeq.map(_.numTasks).sum - val totalSlots = maxConcurrentTasksForStage(stage) - val outstandingForOthers = outstandingTasksForOtherWork(stage, group) + return false + } + val rp = sc.resourceProfileManager.defaultResourceProfile + val demand = pipelinedJobConcurrentTaskDemand(finalRDD, partitions.length) + val totalSlots = maxConcurrentTasksForProfile(rp.id) + // No stage of this job exists yet, so it has no outstanding tasks of its own to exclude; only + // other concurrent jobs' outstanding demand is charged, resource-profile-scoped. + val outstandingForOthers = outstandingTasksForOtherWork(rp.id, Set.empty) val freeSlots = math.max(0, totalSlots - outstandingForOthers) - if (demand > freeSlots) Some((demand, freeSlots)) else None + if (demand > freeSlots) { + logWarning(log"Rejecting job ${MDC(JOB_ID, jobId)}: pipelined stage group needs " + + log"${MDC(NUM_TASKS, demand)} concurrent task slots but only " + + log"${MDC(NUM_SLOTS, freeSlots)} are free") + listener.jobFailed(new SparkException( + errorClass = "CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT", + messageParameters = scala.collection.immutable.Map( + "numTasks" -> demand.toString, "numSlots" -> freeSlots.toString), + cause = null)) + true + } else { + false + } } /** - * The cluster's total concurrent-task capacity for `stage`'s resource profile. Extracted as a + * The cluster's total concurrent-task capacity for the given resource profile. Extracted as a * seam so tests can control it without changing the cluster's core count. Production reads it * from the scheduler backend, exactly as barrier's slot check does. */ - protected def maxConcurrentTasksForStage(stage: Stage): Int = { - val rp = sc.resourceProfileManager.resourceProfileFromId(stage.resourceProfileId) + protected def maxConcurrentTasksForProfile(rpId: Int): Int = { + val rp = sc.resourceProfileManager.resourceProfileFromId(rpId) sc.maxNumConcurrentTasks(rp) } /** - * Outstanding tasks (running plus enqueued), in `stage`'s resource profile, for work OTHER than - * the given pipelined `group` (whose own members must not be charged against the group's own - * admission). Counting enqueued tasks, not just running ones, means a busy neighbor's queued - * backlog is charged against capacity too, so a group is not admitted against slots that other - * work is already committed to using. Resource-profile-scoped to match - * `maxConcurrentTasksForStage`. Extracted as a seam so tests can control occupancy without + * Outstanding tasks (running plus enqueued) in the given resource profile, for work OTHER than + * the stages in `excludeStageIds`. Counting enqueued tasks, not just running ones, means a busy + * neighbor's queued backlog is charged against capacity too, so a group is not admitted against + * slots that other work is already committed to using. Resource-profile-scoped to match + * `maxConcurrentTasksForProfile`. Extracted as a seam so tests can control occupancy without * launching real tasks; returns 0 for a non-TaskSchedulerImpl backend. */ - protected def outstandingTasksForOtherWork(stage: Stage, group: Set[Stage]): Int = + protected def outstandingTasksForOtherWork(rpId: Int, excludeStageIds: Set[Int]): Int = taskScheduler match { case impl: TaskSchedulerImpl => - impl.outstandingTasksForOtherWorkInProfile(stage.resourceProfileId, group.map(_.id)) + impl.outstandingTasksForOtherWorkInProfile(rpId, excludeStageIds) case _ => 0 } @@ -1761,6 +1781,31 @@ private[spark] class DAGScheduler( return } + // v1 (M1, real-time-mode scope) supports a job that is either ALL-regular or ALL-pipelined, + // not a mix. A job whose shuffle graph combines a pipelined shuffle with a regular one is + // rejected up front (before any stage is created). Restricting to a single kind makes an + // all-pipelined job's whole stage graph one pipelined group with no regular prefix, so gang + // admission is decided up front (see the slot check below) rather than mid-DAG once a producer + // is already running. Mixed / mid-DAG groups are a later milestone. + val (hasPipelined, hasRegular) = classifyJobShuffleKinds(finalRDD) + if (hasPipelined && hasRegular) { + logWarning(log"Rejecting job ${MDC(JOB_ID, jobId)}: a job mixing a pipelined shuffle with " + + log"a regular shuffle is not supported") + listener.jobFailed(new SparkException( + "A job that mixes a pipelined shuffle dependency with a regular shuffle dependency is " + + "not supported: v1 requires a job to be either all-regular or all-pipelined.")) + return + } + + // Gang admission for an all-pipelined job: the whole stage graph is one pipelined group, so + // check up front (before any stage is created) that the cluster can run the entire group + // concurrently. If it cannot fit, fail the job now -- no partial scheduler state, and no member + // ever left running while a sibling waits on slots (true all-or-nothing gang admission). Inert + // for a regular job (no pipelined dependency). + if (hasPipelined && rejectUnadmittablePipelinedGroup(jobId, finalRDD, partitions, listener)) { + return + } + var finalStage: ResultStage = null try { // New stage creation may throw an exception if, for example, jobs are run on a @@ -1926,12 +1971,13 @@ private[spark] class DAGScheduler( submitStage(parent) } - // Submitting a parent can abort the job during this recursion: for a pipelined chain - // A->B->C, submitStage(C) recurses into submitStage(B) whose group slot-check covers the - // whole component {A,B,C} and may abortStage(B), which cleans up A/B/C (including this - // stage) and fails the job. If that happened, `stage` is no longer registered; do not - // proceed to co-schedule or re-park it (re-parking would re-insert a job-less stage into - // waitingStages -- a scheduler-state leak). Inert for a non-aborting submit. + // Submitting a parent can abort the job during this recursion (e.g. a parent that has + // exhausted its stage attempts hits abortStage, which cleans up all of the job's stages + // including this one and fails the job). If that happened, `stage` is no longer + // registered; do not proceed to co-schedule or re-park it (re-parking would re-insert a + // job-less stage into waitingStages -- a scheduler-state leak). Inert for a + // non-aborting submit. (Pipelined group admission is decided up front in + // handleJobSubmitted, so it cannot abort the group from within this recursion.) if (!stageIdToStage.contains(stage.id)) { logInfo(log"${MDC(STAGE, stage)} was removed during parent submission (its job was " + log"aborted); not co-scheduling or re-parking it") @@ -1956,32 +2002,22 @@ private[spark] class DAGScheduler( pipelinedMissing.nonEmpty && pipelinedMissing.forall(runningStages.contains) if (regularMissing.isEmpty && allPipelinedParentsRunning) { - // Best-effort slot check: the whole pipelined group must run at once, so its demand - // must fit in the currently-FREE slots (total capacity minus what other work has - // outstanding -- running plus enqueued; spec S4.1). If it cannot, fail fast with a - // clear error rather than deadlock (there is no out-of-band slot reservation in v1). - pipelinedGroupExceedsCapacity(stage) match { - case Some((demand, freeSlots)) => - abortStage(stage, s"Cannot co-schedule pipelined stage group: needs $demand " + - s"concurrent task slots but only $freeSlots are currently free.", - Some(new SparkException( - errorClass = "CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT", - messageParameters = scala.collection.immutable.Map( - "numTasks" -> demand.toString, "numSlots" -> freeSlots.toString), - cause = null))) - case None => - 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, - // so its successful completions are deferred until those producers finish (S5). - val deferral = - pipelinedConsumerDeferrals.getOrElseUpdate(stage, DeferredCompletion()) - deferral.pendingProducers ++= pipelinedMissing - submitMissingTasks(stage, jobId.get) - // This stage is now running; if it is itself the pipelined producer of a waiting - // consumer, co-schedule that consumer too. - submitWaitingPipelinedChildStages(stage) - } + // 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). + 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, + // so its successful completions are deferred until those producers finish (S5). + val deferral = + pipelinedConsumerDeferrals.getOrElseUpdate(stage, DeferredCompletion()) + deferral.pendingProducers ++= pipelinedMissing + submitMissingTasks(stage, jobId.get) + // This stage is now running; if it is itself the pipelined producer of a waiting + // consumer, co-schedule that consumer too. + submitWaitingPipelinedChildStages(stage) } else { waitingStages += stage } 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 32971ca3a140c..694904e714e1f 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -385,14 +385,14 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti // are not gated by the real (local[2]) backend capacity; a test can lower it to exercise the // fail-fast path. @volatile var maxConcurrentTasksForTest: Int = 1000 - override protected def maxConcurrentTasksForStage(stage: Stage): Int = maxConcurrentTasksForTest + override protected def maxConcurrentTasksForProfile(rpId: Int): Int = maxConcurrentTasksForTest // Seam for the free-slot admission check (spec S4.1): outstanding (running + enqueued) task // demand of OTHER work. Default 0 so existing tests see a fully-free cluster; a test can set it // to model a busy neighbor. - @volatile var outstandingTasksForOtherWorkForTest: (Stage, Set[Stage]) => Int = (_, _) => 0 - override protected def outstandingTasksForOtherWork(stage: Stage, group: Set[Stage]): Int = - outstandingTasksForOtherWorkForTest(stage, group) + @volatile var outstandingTasksForOtherWorkForTest: (Int, Set[Int]) => Int = (_, _) => 0 + override protected def outstandingTasksForOtherWork(rpId: Int, excludeStageIds: Set[Int]): Int = + outstandingTasksForOtherWorkForTest(rpId, excludeStageIds) /** * Schedules shuffle merge finalize. @@ -6214,40 +6214,27 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } - test("pipelined shuffle: mixed parents wait for the regular parent, co-schedule the pipelined") { - // consumer depends on a pipelined producer AND a regular producer. It must wait for the - // regular one, but when resubmitted it co-schedules with the (still-running) pipelined one. + test("pipelined shuffle: a job mixing a pipelined and a regular shuffle is rejected up front") { + // v1 (M1) supports a job that is either all-regular or all-pipelined, not a mix. A consumer + // depending on BOTH a pipelined producer AND a regular producer is a mixed job and must be + // rejected up front (before any stage is submitted), leaving no scheduler state behind. val pipelinedProducerRdd = new MyRDD(sc, 2, Nil) val pipelinedDep = new PipelinedShuffleDependency(pipelinedProducerRdd, new HashPartitioner(2)) val regularProducerRdd = new MyRDD(sc, 2, Nil) val regularDep = new ShuffleDependency(regularProducerRdd, new HashPartitioner(2)) val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep, regularDep), tracker = mapOutputTracker) - submit(consumerRdd, Array(0, 1)) + 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) - // Both producers are submitted, but the consumer waits (it has a regular missing parent). - assert(taskSets.size === 2, s"expected both producers submitted, got ${taskSets.size}") - assert(scheduler.waitingStages.exists(_.rdd eq consumerRdd), - "consumer should be waiting on its regular parent") - - // Complete the regular producer. Now the consumer is resubmitted and, since its only remaining - // missing parent is pipelined, co-scheduled with it (a 3rd task set appears). - val regularStageId = taskSets.find { ts => - scheduler.stageIdToStage(ts.stageId).rdd eq regularProducerRdd - }.get.stageId - completeShuffleMapStageSuccessfully(regularStageId, 0, 2) - assert(taskSets.size === 3, "consumer should be co-scheduled once the regular parent finishes") - - // Finish the pipelined producer and the consumer. - val pipelinedStageId = taskSets.find { ts => - scheduler.stageIdToStage(ts.stageId).rdd eq pipelinedProducerRdd - }.get.stageId - completeShuffleMapStageSuccessfully(pipelinedStageId, 0, 2) - 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)) + assert(failure.get() != null, "a mixed pipelined+regular job 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() } @@ -6278,175 +6265,6 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } - test("pipelined shuffle: a pipelined producer behind a regular shuffle is NOT co-scheduled " + - "early") { - // A --regular--> B --pipelined--> C. B cannot run until A materializes, so C must NOT be - // co-scheduled with B while B is still parked. (Regression: the naive check co-scheduled C - // against a not-yet-running B, stranding C -- a hang under slot pressure.) - val rddA = new MyRDD(sc, 2, Nil) - val regAB = new ShuffleDependency(rddA, new HashPartitioner(2)) - val rddB = new MyRDD(sc, 2, List(regAB), tracker = mapOutputTracker) - val psdBC = new PipelinedShuffleDependency(rddB, new HashPartitioner(2)) - val rddC = new MyRDD(sc, 2, List(psdBC), tracker = mapOutputTracker) - submit(rddC, Array(0, 1)) - - // Only A runs initially. B waits on A; C waits on B (its pipelined parent isn't running yet). - assert(taskSets.size === 1, s"expected only A submitted, got ${taskSets.size}") - assert(scheduler.waitingStages.exists(_.rdd eq rddB), "B should wait on its regular parent A") - assert(scheduler.waitingStages.exists(_.rdd eq rddC), - "C must NOT be co-scheduled while its pipelined producer B is still parked") - - // A finishes -> B becomes runnable and is now co-scheduled with C (its pipelined consumer). - val idA = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddA).get.stageId - completeShuffleMapStageSuccessfully(idA, 0, 2) - assert(scheduler.runningStages.exists(_.rdd eq rddB), "B should now be running") - assert(scheduler.runningStages.exists(_.rdd eq rddC), - "C should now be co-scheduled with the running B") - - val idB = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddB).get.stageId - completeShuffleMapStageSuccessfully(idB, 0, 2) - val tsC = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddC).get - complete(tsC, Seq((Success, 42), (Success, 43))) - assert(results === Map(0 -> 42, 1 -> 43)) - assertDataStructuresEmpty() - } - - test("pipelined shuffle: consumer with two pipelined parents co-schedules with both") { - val rddA = new MyRDD(sc, 2, Nil) - val psdA = new PipelinedShuffleDependency(rddA, new HashPartitioner(2)) - val rddB = new MyRDD(sc, 2, Nil) - val psdB = new PipelinedShuffleDependency(rddB, new HashPartitioner(2)) - val consumerRdd = new MyRDD(sc, 2, List(psdA, psdB), tracker = mapOutputTracker) - submit(consumerRdd, Array(0, 1)) - - // Both producers and the consumer are all co-scheduled. - assert(taskSets.size === 3, s"expected both producers + consumer, got ${taskSets.size}") - assert(!scheduler.waitingStages.exists(_.rdd eq consumerRdd)) - Seq(rddA, rddB, consumerRdd).foreach { rdd => - assert(scheduler.runningStages.exists(_.rdd eq rdd)) - } - - val idA = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddA).get.stageId - completeShuffleMapStageSuccessfully(idA, 0, 2) - val idB = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddB).get.stageId - completeShuffleMapStageSuccessfully(idB, 0, 2) - val tsC = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd).get - complete(tsC, Seq((Success, 42), (Success, 43))) - assert(results === Map(0 -> 42, 1 -> 43)) - assertDataStructuresEmpty() - } - - test("pipelined shuffle: consumer with two pipelined parents re-parks until BOTH are running") { - // R1 --regular--> P1 --pipelined--> C ; R2 --regular--> P2 --pipelined--> C. - // When P1 starts (R1 done) but P2 is still parked (R2 not done), C must RE-PARK, not - // co-schedule -- else it would run against a not-yet-running P2. (Guards the `forall` in the - // co-schedule gate against a `.exists` regression.) - val rddR1 = new MyRDD(sc, 2, Nil) - val rddP1 = new MyRDD(sc, 2, List(new ShuffleDependency(rddR1, new HashPartitioner(2))), - tracker = mapOutputTracker) - val psdP1 = new PipelinedShuffleDependency(rddP1, new HashPartitioner(2)) - val rddR2 = new MyRDD(sc, 2, Nil) - val rddP2 = new MyRDD(sc, 2, List(new ShuffleDependency(rddR2, new HashPartitioner(2))), - tracker = mapOutputTracker) - val psdP2 = new PipelinedShuffleDependency(rddP2, new HashPartitioner(2)) - val rddC = new MyRDD(sc, 2, List(psdP1, psdP2), tracker = mapOutputTracker) - submit(rddC, Array(0, 1)) - - // Initially only the two regular roots R1, R2 run; P1, P2, C all wait. - assert(taskSets.size === 2, s"expected R1, R2 submitted, got ${taskSets.size}") - assert(scheduler.waitingStages.exists(_.rdd eq rddC)) - - // R1 completes -> P1 runs and reconsiders C. But P2 is still parked, so C must re-park. - val idR1 = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddR1).get.stageId - completeShuffleMapStageSuccessfully(idR1, 0, 2) - assert(scheduler.runningStages.exists(_.rdd eq rddP1), "P1 should be running") - assert(!scheduler.runningStages.exists(_.rdd eq rddP2), "P2 should not be running yet") - assert(scheduler.waitingStages.exists(_.rdd eq rddC), - "C must re-park while its second pipelined parent P2 is not yet running") - - // R2 completes -> P2 runs and reconsiders C; now both pipelined parents run, so C co-schedules. - val idR2 = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddR2).get.stageId - completeShuffleMapStageSuccessfully(idR2, 0, 2) - assert(scheduler.runningStages.exists(_.rdd eq rddC), "C should now be co-scheduled") - - val idP1 = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddP1).get.stageId - completeShuffleMapStageSuccessfully(idP1, 0, 2) - val idP2 = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddP2).get.stageId - completeShuffleMapStageSuccessfully(idP2, 0, 2) - val tsC = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddC).get - complete(tsC, Seq((Success, 42), (Success, 43))) - assert(results === Map(0 -> 42, 1 -> 43)) - assertDataStructuresEmpty() - } - - // Note: there is deliberately no "one producer feeding two pipelined consumers" test. That shape - // (a single producer stage with two pipelined consumers) is pipelined *fan-out*, which requires a - // shared PipelinedShuffleDependency and is forbidden by the spec (S9); its admission-time - // rejection is a later milestone. Two distinct PipelinedShuffleDependency objects over one RDD do - // NOT create it -- they create two separate producer stages, each with one consumer. So for every - // supported (non-fan-out) shape a running pipelined producer has at most one waiting pipelined - // consumer, and submitWaitingPipelinedChildStages reconsiders it via the single-child path - // covered by the transitive-cascade and producer-behind-regular-shuffle tests. - - test("pipelined shuffle: reconsideration cascades transitively (B->C->D behind a regular root)") { - // Z --regular--> B --pipelined--> C --pipelined--> D. When Z completes, B runs and reconsiders - // C; co-scheduling C must in turn reconsider D (transitive cascade through the co-schedule - // branch's submitWaitingPipelinedChildStages call). - val rddZ = new MyRDD(sc, 2, Nil) - val rddB = new MyRDD(sc, 2, List(new ShuffleDependency(rddZ, new HashPartitioner(2))), - tracker = mapOutputTracker) - val rddC = new MyRDD(sc, 2, List(new PipelinedShuffleDependency(rddB, new HashPartitioner(2))), - tracker = mapOutputTracker) - val rddD = new MyRDD(sc, 2, List(new PipelinedShuffleDependency(rddC, new HashPartitioner(2))), - tracker = mapOutputTracker) - submit(rddD, Array(0, 1)) - - assert(taskSets.size === 1, s"expected only Z submitted, got ${taskSets.size}") - - // Z completes -> B runs -> reconsiders C -> co-schedules C -> reconsiders D -> co-schedules D. - val idZ = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddZ).get.stageId - completeShuffleMapStageSuccessfully(idZ, 0, 2) - assert(scheduler.runningStages.exists(_.rdd eq rddB), "B running") - assert(scheduler.runningStages.exists(_.rdd eq rddC), "C co-scheduled with B") - assert(scheduler.runningStages.exists(_.rdd eq rddD), "D co-scheduled transitively with C") - - val idB = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddB).get.stageId - completeShuffleMapStageSuccessfully(idB, 0, 2) - val idC = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddC).get.stageId - completeShuffleMapStageSuccessfully(idC, 0, 2) - val tsD = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddD).get - complete(tsD, Seq((Success, 42), (Success, 43))) - assert(results === Map(0 -> 42, 1 -> 43)) - assertDataStructuresEmpty() - } - - test("pipelined shuffle: reconsidered consumer is not submitted twice when producer completes") { - // R --regular--> B --pipelined--> C. C is co-scheduled when B starts (reconsideration). When B - // later COMPLETES, submitWaitingChildStages(B) must NOT submit C a second time. - val rddR = new MyRDD(sc, 2, Nil) - val rddB = new MyRDD(sc, 2, List(new ShuffleDependency(rddR, new HashPartitioner(2))), - tracker = mapOutputTracker) - val rddC = new MyRDD(sc, 2, List(new PipelinedShuffleDependency(rddB, new HashPartitioner(2))), - tracker = mapOutputTracker) - submit(rddC, Array(0, 1)) - - val idR = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddR).get.stageId - completeShuffleMapStageSuccessfully(idR, 0, 2) - // C co-scheduled exactly once via reconsideration. - def cTaskSets: Int = taskSets.count(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddC) - assert(cTaskSets === 1, "C should be submitted once by reconsideration") - - // B completes -> submitWaitingChildStages(B) fires, but C is already running, not waiting. - val idB = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddB).get.stageId - completeShuffleMapStageSuccessfully(idB, 0, 2) - assert(cTaskSets === 1, "C must not be submitted a second time when B completes") - - val tsC = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddC).get - complete(tsC, Seq((Success, 42), (Success, 43))) - assert(results === Map(0 -> 42, 1 -> 43)) - 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. @@ -6582,6 +6400,29 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti // Gang admission / slot check (spec S4, S4.1) // ========================================================================================== + test("pipelined shuffle: an all-pipelined group that fits is admitted up front and runs") { + // Whole-group demand producer(2) + consumer(2) = 4 <= capacity 4, other-work occupancy 0, so + // the up-front gang admission check (handleJobSubmitted) admits the job before any stage runs, + // and both stages are then co-scheduled and complete normally. + val producerRdd = new MyRDD(sc, 2, Nil) + val myScheduler = scheduler.asInstanceOf[MyDAGScheduler] + myScheduler.maxConcurrentTasksForTest = 4 + myScheduler.outstandingTasksForOtherWorkForTest = (_, _) => 0 + try { + 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, "a group that fits must be admitted and co-scheduled") + completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2) + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } finally { + myScheduler.maxConcurrentTasksForTest = 1000 + myScheduler.outstandingTasksForOtherWorkForTest = (_, _) => 0 + } + } + test("pipelined shuffle: a group too large to co-fit fails fast with INSUFFICIENT_SLOT") { // Constrain reported capacity to 3 slots. A pipelined group of producer(2) + consumer(2) = 4 // tasks exceeds it and can never be co-resident, so it must fail fast rather than deadlock. @@ -6607,18 +6448,11 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } } - test("pipelined shuffle: an over-capacity 3-stage chain aborts without leaking a stage in " + - "waitingStages (nested-frame abort)") { - // Regression for a nested-frame abort: for a pipelined chain A->B->C, submitStage(C) recurses - // into submitStage(B), whose group slot-check covers the whole component {A,B,C} and aborts the - // job from that nested frame (cleaning up A/B/C). Control then unwinds to the outer - // submitStage(C) frame; without the "was I removed during recursion?" guard, its else branch - // would re-park the already-removed C into waitingStages -- a job-less scheduler-state leak. - // The - // 2-stage over-capacity test does not exercise this (there the abort fires in the consumer's - // own - // frame). Assert the abort leaves NO stage behind (assertDataStructuresEmpty checks - // waitingStages). + test("pipelined shuffle: an over-capacity 3-stage all-pipelined chain is rejected up front") { + // A 3-stage all-pipelined chain A->B->C has whole-group demand 2+2+2 = 6. With capacity 3 it + // cannot co-fit, so the up-front gang admission check (handleJobSubmitted) fails the job before + // any stage is created -- leaving no scheduler state behind (assertDataStructuresEmpty). This + // covers a multi-stage group in addition to the 2-stage "group too large" case. val rddA = new MyRDD(sc, 2, Nil) // touch sc first so `scheduler` is initialized scheduler.asInstanceOf[MyDAGScheduler].maxConcurrentTasksForTest = 3 try { @@ -6637,7 +6471,8 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assert(failure.get().getMessage.contains("CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT") || failure.get().getMessage.contains("concurrent task slots"), s"expected an insufficient-slot error, got: ${failure.get().getMessage}") - // The key assertion: no stage (esp. the outer consumer C) is left behind in waitingStages. + // No stage is created for a job rejected up front. + assert(taskSets.isEmpty, "no stage should be submitted for an up-front-rejected group") assertDataStructuresEmpty() } finally { scheduler.asInstanceOf[MyDAGScheduler].maxConcurrentTasksForTest = 1000 From ede07d22aedc26720bbcfb5b862b3b759ce5c07c Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sun, 19 Jul 2026 06:07:31 +0000 Subject: [PATCH 10/54] [SPARK-XXXXX][CORE] Align pipelined deferred-completion naming with ConcurrentStageDAGScheduler Rename the deferred-completion machinery to match the production RTM scheduler (ConcurrentStageDAGScheduler), so reviewers familiar with that code read the same concepts here. Pure rename; no behavior change. - pipelinedConsumerDeferrals -> dependentStageMap - DeferredCompletion -> DependentStageInfo - pendingProducers -> parents - bufferedEvents -> delayedTaskCompletionEvents The pipelined-specific helpers (isPipelinedProducer, isPipelinedGroupMember, submitWaitingPipelinedChildStages, rejectUnadmittablePipelinedGroup) keep their names -- they belong to the type-driven pipelined model and have no clean analog in the reference scheduler. Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 48 +++++++++---------- .../spark/scheduler/DAGSchedulerSuite.scala | 6 +-- 2 files changed, 27 insertions(+), 27 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 b83dc8ee555a7..8f66537d65919 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -177,15 +177,15 @@ private[spark] class DAGScheduler( * S5, group-observable completion). We buffer such events here and replay them once every * pipelined producer the consumer was waiting on has finished. * - * Keyed by the consumer stage. `pendingProducers` is the set of its pipelined producer stages not - * yet finished; `bufferedEvents` are its Success CompletionEvents held until then. Entries exist - * only for stages that were co-scheduled with a not-yet-finished pipelined producer, so this is - * empty for any job without a pipelined dependency. + * Keyed by the consumer stage. `parents` is the set of its pipelined producer stages not yet + * finished; `delayedTaskCompletionEvents` are its Success CompletionEvents held until then. + * Entries exist only for stages that were co-scheduled with a not-yet-finished pipelined + * producer, so this is empty for any job without a pipelined dependency. */ - private[scheduler] case class DeferredCompletion( - pendingProducers: HashSet[Stage] = new HashSet[Stage], - bufferedEvents: ListBuffer[CompletionEvent] = new ListBuffer[CompletionEvent]) - private[scheduler] val pipelinedConsumerDeferrals = new HashMap[Stage, DeferredCompletion] + private[scheduler] case class DependentStageInfo( + parents: HashSet[Stage] = new HashSet[Stage], + delayedTaskCompletionEvents: ListBuffer[CompletionEvent] = new ListBuffer[CompletionEvent]) + private[scheduler] val dependentStageMap = new HashMap[Stage, DependentStageInfo] private[scheduler] val activeJobs = new HashSet[ActiveJob] @@ -1134,8 +1134,8 @@ private[spark] class DAGScheduler( // Drop any pipelined-completion deferral keyed on this stage (as consumer), and // remove it from other consumers' pending-producer sets (as producer), so no // deferral outlives its job (e.g. on job abort before the producers finished). - pipelinedConsumerDeferrals -= stage - pipelinedConsumerDeferrals.values.foreach(_.pendingProducers -= stage) + dependentStageMap -= stage + dependentStageMap.values.foreach(_.parents -= stage) } // data structures based on StageId stageIdToStage -= stageId @@ -2012,8 +2012,8 @@ private[spark] class DAGScheduler( // Record that this stage is co-scheduled with still-running pipelined producers, // so its successful completions are deferred until those producers finish (S5). val deferral = - pipelinedConsumerDeferrals.getOrElseUpdate(stage, DeferredCompletion()) - deferral.pendingProducers ++= pipelinedMissing + dependentStageMap.getOrElseUpdate(stage, DependentStageInfo()) + deferral.parents ++= pipelinedMissing submitMissingTasks(stage, jobId.get) // This stage is now running; if it is itself the pipelined producer of a waiting // consumer, co-schedule that consumer too. @@ -2823,12 +2823,12 @@ private[spark] class DAGScheduler( // the side effects normally. This deferral check must therefore precede updateAccumulators and // postTaskEnd. Inert for jobs with no pipelined dependency (the map is empty). if (event.reason == Success) { - pipelinedConsumerDeferrals.get(stage) match { - case Some(deferral) if deferral.pendingProducers.nonEmpty => + dependentStageMap.get(stage) match { + case Some(deferral) if deferral.parents.nonEmpty => logInfo(log"Deferring completion of task ${MDC(TASK_ID, event.taskInfo.taskId)} in " + log"pipelined consumer ${MDC(STAGE, stage)} until its producer(s) " + - log"${MDC(MISSING_PARENT_STAGES, deferral.pendingProducers.toSeq)} finish") - deferral.bufferedEvents += event + log"${MDC(MISSING_PARENT_STAGES, deferral.parents.toSeq)} finish") + deferral.delayedTaskCompletionEvents += event return case _ => } @@ -3925,19 +3925,19 @@ private[spark] class DAGScheduler( */ private def releaseDeferredPipelinedConsumers( finishedStage: Stage, producerFailed: Boolean): Unit = { - if (pipelinedConsumerDeferrals.isEmpty) { + if (dependentStageMap.isEmpty) { return } // Consumers waiting on this producer. - val released = pipelinedConsumerDeferrals.filter { - case (_, d) => d.pendingProducers.contains(finishedStage) + val released = dependentStageMap.filter { + case (_, d) => d.parents.contains(finishedStage) }.keys.toArray for (consumer <- released) { - val deferral = pipelinedConsumerDeferrals(consumer) - deferral.pendingProducers -= finishedStage - if (deferral.pendingProducers.isEmpty) { - pipelinedConsumerDeferrals -= consumer - val events = deferral.bufferedEvents.toList + val deferral = dependentStageMap(consumer) + deferral.parents -= finishedStage + if (deferral.parents.isEmpty) { + dependentStageMap -= consumer + val events = deferral.delayedTaskCompletionEvents.toList if (producerFailed) { logInfo(log"Dropping ${MDC(NUM_EVENTS, events.size.toLong)} deferred completion(s) for " + log"pipelined consumer ${MDC(STAGE, consumer)} because its producer " + 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 694904e714e1f..b394569134e47 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6712,7 +6712,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) // Deferred while the producer runs. - assert(scheduler.pipelinedConsumerDeferrals.keys.exists(_.rdd eq consumerRdd)) + assert(scheduler.dependentStageMap.keys.exists(_.rdd eq consumerRdd)) assert(results.isEmpty) val producerStage = @@ -6723,7 +6723,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti // (ShuffleMapStage && !producerFailed && !isAvailable) and must NOT release the consumer. assert(!producerStage.isAvailable, "precondition: producer not yet available") scheduler.markStageAsFinished(producerStage, errorMessage = None, willRetry = false) - assert(scheduler.pipelinedConsumerDeferrals.keys.exists(_.rdd eq consumerRdd), + assert(scheduler.dependentStageMap.keys.exists(_.rdd eq consumerRdd), "a not-yet-available pipelined producer must NOT release its deferred consumers (it is " + "about to resubmit; releasing would apply results against soon-to-be-recomputed output)") assert(results.isEmpty, @@ -6735,7 +6735,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti scheduler.runningStages += producerStage completeShuffleMapStageSuccessfully(producerStageId, 0, 2) assert(producerStage.isAvailable) - assert(!scheduler.pipelinedConsumerDeferrals.keys.exists(_.rdd eq consumerRdd), + assert(!scheduler.dependentStageMap.keys.exists(_.rdd eq consumerRdd), "deferral released once the producer is genuinely done (available)") assert(results === Map(0 -> 42, 1 -> 43)) assertDataStructuresEmpty() From 3ea8e80ecb24365e2d7fe36efbc0608c62a27db1 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sun, 19 Jul 2026 19:15:25 +0000 Subject: [PATCH 11/54] [SPARK-XXXXX][CORE] Review round 1: add dependentStageMap to the leak check; fix a vacuous admission test Two test-hygiene fixes from an adversarial review: - assertDataStructuresEmpty now asserts dependentStageMap.isEmpty, per the DAGScheduler class-doc checklist ("when adding a new data structure, update DAGSchedulerSuite.assertDataStructuresEmpty ... to catch memory leaks"). The new dependentStageMap (the pipelined deferred-completion buffer) was omitted; it is always drained today, so this passes -- and now guards a future leak of buffered CompletionEvents / a consumer stage left as perpetually-running. - Retitle/reframe "the group's own running members are not charged against its admission": it stubbed the occupancy seam to a constant ignoring excludeStageIds, and the up-front admission passes Set.empty (no job stage exists yet), so it never exercised member exclusion -- it only asserts a demand==free-capacity group is admitted (a real boundary). The member-exclusion invariant is genuinely covered in TaskSchedulerImplSuite; the title/comment now say what the test actually verifies. Co-authored-by: Isaac --- .../spark/scheduler/DAGSchedulerSuite.scala | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 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 b394569134e47..f8b0e52e81ec0 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6130,6 +6130,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assert(scheduler.runningStages.isEmpty) assert(scheduler.shuffleIdToMapStage.isEmpty) assert(scheduler.waitingStages.isEmpty) + assert(scheduler.dependentStageMap.isEmpty) assert(scheduler.outputCommitCoordinator.isEmpty) } @@ -6552,17 +6553,17 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } } - test("pipelined shuffle: the group's own running members are not charged against its admission") { - // Occupancy for the free-slot check counts only OTHER work: the group's own already-running - // producer must not be charged, or a group that exactly fills free capacity would be wrongly - // rejected. Model a full cluster where the ONLY running tasks are this group's own: 4 total - // slots, other-work occupancy = 0 (the seam excludes the group) -> free = 4 >= demand 4 -> - // admitted. The exclusion itself is unit-tested against TaskSchedulerImpl separately; here we - // assert the DAGScheduler admits a demand==capacity group when nothing else is running. + test("pipelined shuffle: a group whose demand exactly equals free capacity is admitted") { + // Boundary: demand == free capacity must be admitted, not rejected by an off-by-one (the check + // is `demand > freeSlots`, not `>=`). 4 total slots, other-work occupancy 0 -> free 4 >= demand + // 4 -> admitted. NOTE: the DAGScheduler up-front admission (rejectUnadmittablePipelinedGroup) + // passes excludeStageIds = Set.empty (no job stage exists yet at admission time), so the + // group's-own-members exclusion is NOT exercised here; that exclusion is genuinely covered by + // TaskSchedulerImplSuite ("outstandingTasksForOtherWorkInProfile excludes the given stages"). val producerRdd = new MyRDD(sc, 2, Nil) val myScheduler = scheduler.asInstanceOf[MyDAGScheduler] myScheduler.maxConcurrentTasksForTest = 4 - myScheduler.outstandingTasksForOtherWorkForTest = (_, _) => 0 // only the group's own work runs + myScheduler.outstandingTasksForOtherWorkForTest = (_, _) => 0 try { val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) From dfec9c85b593a03be93294c2698a6114782ed7b3 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 06:12:53 +0000 Subject: [PATCH 12/54] [SPARK-XXXXX][CORE] Pipelined consumer deferral: defer only completion, run per-task effects inline (S5.1) The group-observable completion for a pipelined consumer (spec S5) previously used a coarse model: when a consumer finished ahead of its still-running producer, the WHOLE CompletionEvent was buffered and its side effects were withheld until replay. That contradicts the negotiated spec S5.1, which requires per-task side effects -- accumulator updates and the SparkListenerTaskEnd listener event -- to flow in real time, deferring ONLY the stage/job-completion decision (advancing that early is what would cancel the still-running producer or expose the consumer's output prematurely). Switch to the fine-grained model: - handleTaskCompletion runs a deferred consumer's per-task effects inline (as for any other stage), then buffers the event and returns before the completion bookkeeping. - A new CompletionEvent.finishOnly flag marks a replayed event; on replay only the completion bookkeeping runs, so the per-task effects are never applied twice. - releaseDeferredPipelinedConsumers replays with finishOnly=true; the producer-failure drop path simply discards the buffered events (their TaskEnd already fired inline), removing the old asymmetry where the drop path re-emitted postTaskEnd but not updateAccumulators. Observability improves: a finished consumer's tasks are no longer reported as running for the producer's remaining lifetime, and no late TaskEnd burst appears at group abort. Tests: the two deferral tests now assert per-task TaskEnd fires in real time (before the producer finishes) and exactly once. Reverting the fix makes them fail (TaskEnd count 0 instead of 2 while the producer runs). Full DAGSchedulerSuite green. Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 129 ++++++++++-------- .../spark/scheduler/DAGSchedulerEvent.scala | 8 +- .../spark/scheduler/DAGSchedulerSuite.scala | 93 ++++++++----- 3 files changed, 144 insertions(+), 86 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 8f66537d65919..c8193058c4d38 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -2811,58 +2811,69 @@ private[spark] class DAGScheduler( val stage = stageIdToStage(task.stageId) - // Group-observable completion (S5): if this stage is a pipelined consumer co-scheduled with a - // still-running pipelined producer, defer its *successful* completion in full until the - // producer(s) finish. Buffer the whole CompletionEvent and return before ANY of its side - // effects run (accumulator update, task-end listener event, stage/job completion) -- otherwise - // a consumer finishing ahead of its producer would advance job completion and cancel the - // still-running producer, or expose its output early. Deferring the entire event (not just the - // completion bookkeeping) is the coarse model, and crucially it makes the side effects run - // exactly ONCE, at replay: markStageAsFinished re-posts the buffered event once the last - // producer completes (or drops it if a producer fails, S6), and it then re-enters here and runs - // the side effects normally. This deferral check must therefore precede updateAccumulators and - // postTaskEnd. Inert for jobs with no pipelined dependency (the map is empty). - if (event.reason == Success) { - dependentStageMap.get(stage) match { - case Some(deferral) if deferral.parents.nonEmpty => - logInfo(log"Deferring completion of task ${MDC(TASK_ID, event.taskInfo.taskId)} in " + - log"pipelined consumer ${MDC(STAGE, stage)} until its producer(s) " + - log"${MDC(MISSING_PARENT_STAGES, deferral.parents.toSeq)} finish") - deferral.delayedTaskCompletionEvents += event - return + // Group-observable completion (spec S5), fine-grained model: a pipelined consumer's per-task + // side effects (accumulator update, task-end listener event) run in real time as its tasks + // finish -- exactly as for any other stage -- and ONLY the stage/job-completion decision is + // deferred until the producer(s) finish. Deferring the completion decision is what matters: + // advancing it early would let a consumer that finished ahead of its producer complete the job + // and cancel the still-running producer (via cancelRunningIndependentStages), or expose the + // consumer's output before the producer's. Its per-task facts, by contrast, are true when they + // happen (S5.1: task-level listener events flow in real time), so they must not be frozen for + // the producer's remaining lifetime. + // + // Two flags drive this. `finishOnly` is set on a replayed event (see + // releaseDeferredPipelinedConsumers): its per-task effects already ran inline on first + // completion, so on replay we skip straight to the completion bookkeeping and never apply them + // twice. `deferCompletion` is true for a fresh Success from a pipelined consumer whose + // producers are still running: run the per-task effects now, then buffer the event and return + // before the completion bookkeeping. Both are inert for a job with no pipelined dependency (the + // map is empty), so the regular path is unchanged. + val deferCompletion = !event.finishOnly && event.reason == Success && + dependentStageMap.get(stage).exists(_.parents.nonEmpty) + + if (!event.finishOnly) { + // Make sure the task's accumulators are updated before any other processing happens, so that + // we can post a task end event before any jobs or stages are updated. The accumulators are + // only updated in certain cases. + event.reason match { + case Success => + task match { + case rt: ResultTask[_, _] => + val resultStage = stage.asInstanceOf[ResultStage] + resultStage.activeJob match { + case Some(job) => + // Only update the accumulator once for each result task. + if (!job.finished(rt.outputId)) { + updateAccumulators(event) + } + case None => // Ignore update if task's job has finished. + } + case _ => + updateAccumulators(event) + } + case _: ExceptionFailure | _: TaskKilled => updateAccumulators(event) case _ => } - } + if (trackingCacheVisibility) { + // Update rdd blocks' visibility status. + blockManagerMaster.updateRDDBlockVisibility( + event.taskInfo.taskId, visible = event.reason == Success) + } - // Make sure the task's accumulators are updated before any other processing happens, so that - // we can post a task end event before any jobs or stages are updated. The accumulators are - // only updated in certain cases. - event.reason match { - case Success => - task match { - case rt: ResultTask[_, _] => - val resultStage = stage.asInstanceOf[ResultStage] - resultStage.activeJob match { - case Some(job) => - // Only update the accumulator once for each result task. - if (!job.finished(rt.outputId)) { - updateAccumulators(event) - } - case None => // Ignore update if task's job has finished. - } - case _ => - updateAccumulators(event) - } - case _: ExceptionFailure | _: TaskKilled => updateAccumulators(event) - case _ => - } - if (trackingCacheVisibility) { - // Update rdd blocks' visibility status. - blockManagerMaster.updateRDDBlockVisibility( - event.taskInfo.taskId, visible = event.reason == Success) + postTaskEnd(event) } - postTaskEnd(event) + // Now that the per-task effects have run, defer the completion bookkeeping if this is a + // co-scheduled pipelined consumer whose producer(s) are still running. It is replayed + // (finishOnly = true) once the last producer finishes, or dropped if a producer fails (S6). + if (deferCompletion) { + val deferral = dependentStageMap(stage) + logInfo(log"Deferring completion of task ${MDC(TASK_ID, event.taskInfo.taskId)} in " + + log"pipelined consumer ${MDC(STAGE, stage)} until its producer(s) " + + log"${MDC(MISSING_PARENT_STAGES, deferral.parents.toSeq)} finish") + deferral.delayedTaskCompletionEvents += event + return + } event.reason match { case Success => @@ -3921,7 +3932,15 @@ private[spark] class DAGScheduler( * consumer was deferred on, remove it from that consumer's pending-producer set. When a consumer * has no pending producers left, either replay its buffered completion events (producer succeeded) * or drop them (a producer failed -- the group will be torn down and rerun, so the consumer's - * buffered successes must not be applied; S6). Inert unless `finishedStage` is a tracked producer. + * buffered completions must not be applied; S6). Inert unless `finishedStage` is a tracked + * producer. + * + * Only the deferred stage/job-completion bookkeeping is buffered here; each event's per-task side + * effects (accumulator update, task-end listener event) already ran inline when the task first + * completed (fine-grained model, S5). So replay re-posts each event with `finishOnly = true` (run + * the completion bookkeeping only, never the per-task effects again), and the drop path simply + * discards the buffered events -- there is nothing left to emit, since the consumer's TaskEnd + * events were already delivered in real time. */ private def releaseDeferredPipelinedConsumers( finishedStage: Stage, producerFailed: Boolean): Unit = { @@ -3942,15 +3961,17 @@ private[spark] class DAGScheduler( logInfo(log"Dropping ${MDC(NUM_EVENTS, events.size.toLong)} deferred completion(s) for " + log"pipelined consumer ${MDC(STAGE, consumer)} because its producer " + log"${MDC(STAGE, finishedStage)} failed; the group will be rerun") - // The buffered tasks genuinely succeeded, so still emit their TaskEnd events -- otherwise - // listeners that track active tasks (e.g. AppStatusListener) would believe these tasks - // are still running. We deliberately do NOT run the stage/job completion bookkeeping: - // the group failed and will be rerun, so the consumer's results must not be applied. - events.foreach(postTaskEnd) + // Only the deferred stage/job-completion bookkeeping was buffered; the consumer's + // per-task effects (including its TaskEnd events) already ran inline when the tasks first + // completed. The group failed and will be rerun, so that completion bookkeeping must NOT + // be applied -- simply discard the buffered events. (There is nothing to re-emit: unlike + // the coarse model, no TaskEnd was withheld.) } else { logInfo(log"Replaying ${MDC(NUM_EVENTS, events.size.toLong)} deferred completion(s) for " + log"pipelined consumer ${MDC(STAGE, consumer)} now that its producers have finished") - events.foreach(eventProcessLoop.post) + // Re-post each event to run ONLY the deferred completion bookkeeping (finishOnly = true); + // the per-task effects already ran inline and must not be applied again. + events.foreach(e => eventProcessLoop.post(e.copy(finishOnly = true))) } } } diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGSchedulerEvent.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGSchedulerEvent.scala index cc788e5c65bcc..58481e0517891 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGSchedulerEvent.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGSchedulerEvent.scala @@ -92,7 +92,13 @@ private[scheduler] case class CompletionEvent( result: Any, accumUpdates: Seq[AccumulatorV2[_, _]], metricPeaks: Array[Long], - taskInfo: TaskInfo) + taskInfo: TaskInfo, + // Set only on a pipelined consumer's completion event that was deferred and is now being + // replayed after its producer(s) finished (spec S5, fine-grained model). Its per-task side + // effects (accumulator update, task-end listener event) already ran inline when the task first + // completed; on replay only the deferred stage/job-completion bookkeeping must run, so this + // flag tells handleTaskCompletion to skip the per-task effects and avoid applying them twice. + finishOnly: Boolean = false) extends DAGSchedulerEvent private[scheduler] case class ExecutorAdded(execId: String, host: String) extends DAGSchedulerEvent 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 f8b0e52e81ec0..57ed13aa1798f 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6630,40 +6630,63 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } test("pipelined shuffle: deferred consumer completions are dropped when the producer fails") { - // If a producer fails while a co-scheduled consumer's completions are buffered, those buffered - // successes must be DROPPED (not applied): the group is torn down / the job fails, and applying - // a partial consumer success would be incorrect (S6). - 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) + // If a producer fails while a co-scheduled consumer's completion is deferred, the deferred + // stage/job-completion bookkeeping must be DROPPED (not applied): the group is torn down / the + // job fails, and applying a partial consumer success would be incorrect (S6). The consumer's + // per-task side effects (TaskEnd) already ran inline when its tasks finished (fine-grained + // model, S5.1), so on the drop path there is nothing to re-emit and no TaskEnd is duplicated. + val taskEndCount = new java.util.concurrent.atomic.AtomicInteger(0) + val countingListener = new SparkListener { + override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = taskEndCount.incrementAndGet() } - submit(consumerRdd, Array(0, 1), listener = failListener) - val producerTaskSet = taskSets.head - val consumerTaskSet = taskSets(1) - - // Consumer finishes early -> its completions are buffered (deferred). - complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) - assert(results.isEmpty, "consumer completions should be buffered while the producer runs") + 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 producerTaskSet = taskSets.head + val consumerTaskSet = taskSets(1) - // The producer now FAILS its whole task set. The job fails; the buffered consumer successes - // must NOT have been applied as results. - failed(producerTaskSet, "producer blew up") - assert(failure.get() != null, "job should fail when the producer fails") - assert(results.isEmpty, - "buffered consumer successes must be dropped when the producer fails, not applied") - assertDataStructuresEmpty() + // Consumer finishes early -> its completion bookkeeping is deferred, but its 2 TaskEnd events + // fire in real time. + complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) + assert(results.isEmpty, "consumer job result should be deferred while the producer runs") + sc.listenerBus.waitUntilEmpty(10000) + assert(taskEndCount.get() === 2, + "consumer's TaskEnd events must fire inline, before the producer's fate; got " + + taskEndCount.get()) + + // The producer now FAILS its whole task set. The job fails; the deferred consumer completion + // must NOT be applied as a result, and no consumer TaskEnd is re-emitted on the drop path. + failed(producerTaskSet, "producer blew up") + assert(failure.get() != null, "job should fail when the producer fails") + assert(results.isEmpty, + "deferred consumer completion must be dropped when the producer fails, not applied") + sc.listenerBus.waitUntilEmpty(10000) + assert(taskEndCount.get() === 2, + "drop path must not re-emit consumer TaskEnd events (they already fired); got " + + taskEndCount.get()) + assertDataStructuresEmpty() + } finally { + sc.removeSparkListener(countingListener) + } } - 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 - // update) applied exactly once -- at replay -- not once when buffered and again when replayed. - // The producer (2 tasks) and consumer (2 tasks) yield exactly 4 TaskEnd events total; a - // buffer+replay double-post would inflate that count. + test("pipelined shuffle: a deferred consumer task fires its TaskEnd exactly once, in real time") { + // Fine-grained deferral (spec S5.1): a deferred consumer's per-task side effects (task-end + // listener event, accumulator update) run in REAL TIME as its tasks finish -- NOT withheld + // until replay -- and each runs exactly once (the completion bookkeeping alone is deferred, + // then replayed with finishOnly=true, which must not re-post the TaskEnd). The producer (2 + // tasks) and consumer (2 tasks) yield exactly 4 TaskEnd events total; a buffer+replay + // double-post would inflate that to 6, and a withhold-until-replay bug would delay the 2 + // consumer events. val taskEndCount = new java.util.concurrent.atomic.AtomicInteger(0) val countingListener = new SparkListener { override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = taskEndCount.incrementAndGet() @@ -6677,9 +6700,17 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti val producerStageId = taskSets.head.stageId val consumerTaskSet = taskSets(1) - // Consumer finishes first -> its two completions are buffered (deferred, no TaskEnd yet). + // Consumer finishes first. Its completion BOOKKEEPING is buffered (no result yet), but its + // per-task TaskEnd events fire immediately (fine-grained model). complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) - // Producer finishes -> the two deferred consumer completions replay. + assert(results.isEmpty, "consumer job result must be deferred until the producer finishes") + sc.listenerBus.waitUntilEmpty(10000) + assert(taskEndCount.get() === 2, + "consumer's 2 TaskEnd events must fire in real time, not at replay; got " + + taskEndCount.get()) + + // Producer finishes -> the two deferred consumer completions replay (finishOnly=true): the + // job result is applied now, but NO additional TaskEnd is posted (they already fired above). completeShuffleMapStageSuccessfully(producerStageId, 0, 2) assert(results === Map(0 -> 42, 1 -> 43)) sc.listenerBus.waitUntilEmpty(10000) From 10dd37e6665ed1231117c41ab5114e0d2e9993b6 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 06:23:21 +0000 Subject: [PATCH 13/54] [SPARK-XXXXX][CORE] Fix isPipelinedGroupMember scaladoc: drop stale "reads as unused" and a phantom symbol The NOTE on isPipelinedGroupMember said it "reads as unused when this change is viewed alone" and listed `pipelinedGroupOf` as a sibling helper. Both are inaccurate in the merged view: the method IS used (its call sites -- TaskSet.isPipelined tagging and the member-FetchFailed group abort -- are part of this stack), and `pipelinedGroupOf` does not exist. Reword to name the actual call sites and drop the phantom reference. Co-authored-by: Isaac --- .../scala/org/apache/spark/scheduler/DAGScheduler.scala | 8 ++++---- 1 file changed, 4 insertions(+), 4 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 c8193058c4d38..a1e931add43e7 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -1604,10 +1604,10 @@ private[spark] class DAGScheduler( * fail-fast behavior off this (via TaskSet.isPipelined). False for any stage in a job with no * pipelined dependency. * - * NOTE: this is deliberately defined here alongside the other group-topology helpers - * (`isPipelinedProducer`, `pipelinedGroupOf`) but is first consumed by the group-atomic failure - * handling added in a later PR of this stack (`TaskSet.isPipelined` and the member-failure - * fail-fast, spec S6). It reads as unused when this change is viewed alone; that is expected. + * NOTE: this is deliberately defined here alongside the other group-topology helper + * (`isPipelinedProducer`). Its call sites are the group-atomic failure handling added later in + * this stack: tagging the member's `TaskSet.isPipelined` at submission and routing a member + * FetchFailed to a whole-group abort (spec S6). */ private def isPipelinedGroupMember(stage: Stage): Boolean = { if (isPipelinedProducer(stage)) { From e73d4130efb106d671ca9fec0eb325adda380de5 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 06:57:21 +0000 Subject: [PATCH 14/54] [SPARK-XXXXX][CORE] Tighten outstandingTasksForOtherWorkInProfile to private[scheduler] The method exists only to back the pipelined-group slot admission check in DAGScheduler; it does not belong on the public TaskSchedulerImpl surface. Its only callers (DAGScheduler and TaskSchedulerImplSuite) are in org.apache.spark.scheduler, so private[scheduler] is sufficient and keeps the new capacity-accounting method out of the public API. Co-authored-by: Isaac --- .../scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala index e162a5c19cb69..b832bc6463f16 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala @@ -174,7 +174,7 @@ private[spark] class TaskSchedulerImpl( * count is one consistent snapshot (both the per-profile total and the excluded members are read * together). */ - def outstandingTasksForOtherWorkInProfile( + private[scheduler] def outstandingTasksForOtherWorkInProfile( resourceProfileId: Int, excludeStageIds: Set[Int]): Int = synchronized { taskSetsByStageIdAndAttempt.iterator.flatMap { case (stageId, attempts) => if (excludeStageIds.contains(stageId)) { From c9f2bb091d1f8506f7302122110c307dcac00bec Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 06:59:45 +0000 Subject: [PATCH 15/54] [SPARK-XXXXX][CORE] Clarify pipelined admission vs. barrier: contrast the retry behavior The gang-admission scaladoc likened the check to barrier's slot check. The likeness is only that both reject before any stage is created and compute demand from the RDD graph; their retry behavior differs and a reader should not infer equivalence. Spell out the contrast: barrier re-posts the job and retries its check up to a max-failures bound, whereas pipelined admission is terminal (one check, then fail) and delegates transient-shortfall retry to the caller -- scheduler-side PG-admission retry being a post-v1 refinement (spec S4.1). Co-authored-by: Isaac --- .../scala/org/apache/spark/scheduler/DAGScheduler.scala | 7 +++++++ 1 file changed, 7 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 a1e931add43e7..be0cb9eb95f89 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -1540,6 +1540,13 @@ private[spark] class DAGScheduler( * `listener` with no retry -- a transient shortfall is the caller's to retry (e.g. the streaming * batch loop reruns the batch). Returns true (job failed) if it does not fit. * + * The likeness to barrier's slot check is only that both reject before any stage is created + * (leaving no partial state) and compute demand from the RDD graph. Retry behavior differs + * deliberately: barrier RE-POSTS the job and re-runs its check on a timer up to + * `spark.scheduler.barrier.maxConcurrentTasksCheck.maxFailures` times; pipelined admission is + * TERMINAL (one check, then fail), delegating transient-shortfall retry to the caller + * (scheduler-side PG-admission retry is a post-v1 refinement; spec S4.1). + * * The check can be turned off with `spark.scheduler.pipelinedGroup.slotCheck.enabled=false` (for * deployments that admit capacity out-of-band, e.g. via a slot reservation). */ From fc10be3cd94427bdff164e6ac7b6756f95784f8d Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 22:14:48 +0000 Subject: [PATCH 16/54] [SPARK-XXXXX][CORE] Don't re-post TaskEnd for a finishOnly replay that lands on a torn-down stage A deferred pipelined-consumer completion fires its TaskEnd inline on first completion; only the stage/job bookkeeping is re-posted (finishOnly=true) once the producer finishes. The cancelled-stage guard in handleTaskCompletion (which posts TaskEnd for an event whose stage is no longer tracked) ran unconditionally and before the finishOnly check. If the group is torn down (a sibling aborts) between a finishOnly re-post and its dequeue, the replay lands on the removed stage and the guard would re-post TaskEnd -- double-counting it in listeners (AppStatusListener's active-task count could go negative). Skip postTaskEnd for a finishOnly event in that guard: its per-task TaskEnd already fired inline. A regression test on the failure PR (which has the group-atomic abort this needs) drives a finishOnly replay onto a torn-down stage and asserts no extra TaskEnd; reverting this guard makes it fail (TaskEnd count goes up by one). Co-authored-by: Isaac --- .../scala/org/apache/spark/scheduler/DAGScheduler.scala | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) 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 be0cb9eb95f89..8d74a68c35904 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -2810,7 +2810,14 @@ private[spark] class DAGScheduler( // are properly notified and can chose to handle it. For instance, some listeners are // doing their own accounting and if they don't get the task end event they think // tasks are still running when they really aren't. - postTaskEnd(event) + // + // Exception: a replayed pipelined-consumer completion (finishOnly) already fired its TaskEnd + // inline when the task first completed; only the deferred stage/job bookkeeping was + // re-posted. If the stage was torn down (e.g. a sibling aborted the group) before this replay + // dequeues, re-posting TaskEnd here would double-count it in listeners. Skip for finishOnly. + if (!event.finishOnly) { + postTaskEnd(event) + } // Skip all the actions if the stage has been cancelled. return From fd39ce392d9a9f392e32b6a6ff4c1181b74810e5 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 23:38:27 +0000 Subject: [PATCH 17/54] [SPARK-XXXXX][CORE] Clarify the deferral drop-path comment: inline accumulator deltas are not un-merged The drop-path comment implied full cleanliness ("only the deferred completion bookkeeping was buffered ... must NOT be applied"). But a deferred consumer's per-task effects, including updateAccumulators, run inline before buffering, so on a producer-failure drop the already- applied accumulator deltas are NOT un-merged. Note that this is consistent with base Spark (accumulators are non-transactional across abort+rerun): the rerun re-delivers under the S5 idempotent-sink contract, and RTM SQL metrics are fresh per batch, so there is no cross-batch double-count. Comment-only. Co-authored-by: Isaac --- .../scala/org/apache/spark/scheduler/DAGScheduler.scala | 8 ++++++-- 1 file changed, 6 insertions(+), 2 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 8d74a68c35904..a40b3d2420d42 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -3976,10 +3976,14 @@ private[spark] class DAGScheduler( log"pipelined consumer ${MDC(STAGE, consumer)} because its producer " + log"${MDC(STAGE, finishedStage)} failed; the group will be rerun") // Only the deferred stage/job-completion bookkeeping was buffered; the consumer's - // per-task effects (including its TaskEnd events) already ran inline when the tasks first + // per-task effects (TaskEnd, accumulator updates) already ran inline when the tasks first // completed. The group failed and will be rerun, so that completion bookkeeping must NOT // be applied -- simply discard the buffered events. (There is nothing to re-emit: unlike - // the coarse model, no TaskEnd was withheld.) + // the coarse model, no TaskEnd was withheld.) The already-applied inline accumulator + // deltas are NOT un-merged here -- consistent with base Spark, whose accumulators are + // non-transactional across an abort+rerun; the rerun re-delivers under the S5 idempotent- + // sink contract, and RTM SQL metrics are fresh per batch so there is no cross-batch + // double-count. } else { logInfo(log"Replaying ${MDC(NUM_EVENTS, events.size.toLong)} deferred completion(s) for " + log"pipelined consumer ${MDC(STAGE, consumer)} now that its producers have finished") From 89b29a816d72d5de1d744194defd77ea4669917e Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Tue, 21 Jul 2026 07:16:51 +0000 Subject: [PATCH 18/54] [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, and fine-grained/coarse model terminology -- from the pipelined-shuffle comments, one test title, and the mixed-job error message, rewording to plain behavioral language. Comment, test-title, and error/assertion-string text only; no behavioral change. Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 131 +++++++++--------- .../spark/scheduler/DAGSchedulerEvent.scala | 8 +- .../spark/scheduler/DAGSchedulerSuite.scala | 26 ++-- 3 files changed, 80 insertions(+), 85 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 a40b3d2420d42..34bbae09e35a3 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -171,11 +171,11 @@ private[spark] class DAGScheduler( /** * Deferred completion for pipelined groups. When a stage is co-scheduled with a pipelined - * producer that is still running (a "pipelined consumer"), its successful task-completion events - * must not be processed yet: doing so could finish the consumer's job and cancel the - * still-running producer, or make the consumer's output observable before the producer's (spec - * S5, group-observable completion). We buffer such events here and replay them once every - * pipelined producer the consumer was waiting on has finished. + * producer that is still running (a "pipelined consumer"), the stage/job-completion decision from + * its successful task-completion events must not be processed yet: doing so could finish the + * consumer's job and cancel the still-running producer, or make the consumer's output observable + * before the producer's. We buffer such events here and replay them once every pipelined producer + * the consumer was waiting on has finished. * * Keyed by the consumer stage. `parents` is the set of its pipelined producer stages not yet * finished; `delayedTaskCompletionEvents` are its Success CompletionEvents held until then. @@ -990,10 +990,10 @@ private[spark] class DAGScheduler( * contains any [[PipelinedShuffleDependency]] and whether it contains any regular (non-pipelined) * `ShuffleDependency`. Narrow dependencies are not boundaries and are ignored. * - * v1 (M1, real-time-mode scope) supports a job that is either ALL-regular or ALL-pipelined: a - * job whose shuffle graph mixes a pipelined shuffle with a regular one is rejected fail-fast - * (see `handleJobSubmitted`). Restricting to a single kind makes the whole job one pipelined - * group when pipelined, so gang admission can be decided up front before any stage is submitted. + * A job must be either ALL-regular or ALL-pipelined: a job whose shuffle graph mixes a pipelined + * shuffle with a regular one is rejected fail-fast (see `handleJobSubmitted`). Restricting to a + * single kind makes the whole job one pipelined group when pipelined, so gang admission can be + * decided up front before any stage is submitted. */ private def classifyJobShuffleKinds(finalRDD: RDD[_]): (Boolean, Boolean) = { var hasPipelined = false @@ -1024,7 +1024,7 @@ private[spark] class DAGScheduler( * * Rejected combinations: * - Speculation: a speculative copy of a producer would race a consumer already reading the - * producer's partial output, with no commit barrier protecting the read (spec S9). + * producer's partial output, with no commit barrier protecting the read. * - Dynamic allocation: a pipelined group is gang-scheduled (all member stages must run at * once), and the free-slot admission check measures currently-active executors. Under dynamic * allocation a job can start before any executor has spun up, so the group would be failed @@ -1503,12 +1503,12 @@ private[spark] class DAGScheduler( /** * The total concurrent-task demand of an all-pipelined job, computed from the RDD graph BEFORE * any stage is created (so a rejection based on it leaves no partial scheduler state, exactly as - * the barrier slot check and the speculation/DA reject do). Because a v1 job is either - * all-regular or all-pipelined, an all-pipelined job's whole stage graph is one pipelined group; - * its members are the final result stage plus every pipelined producer. Each member's task count - * is its RDD's partition count (`rdd.partitions.length`), matching how `createShuffleMapStage` - * derives `numTasks`. `finalNumPartitions` is the result stage's task count (the number of - * partitions the job runs, which may be a subset of `finalRDD.partitions`). + * the barrier slot check and the speculation/DA reject do). Because a job is either all-regular + * or all-pipelined, an all-pipelined job's whole stage graph is one pipelined group; its members + * are the final result stage plus every pipelined producer. Each member's task count is its RDD's + * partition count (`rdd.partitions.length`), matching how `createShuffleMapStage` derives + * `numTasks`. `finalNumPartitions` is the result stage's task count (the number of partitions the + * job runs, which may be a subset of `finalRDD.partitions`). */ private def pipelinedJobConcurrentTaskDemand(finalRDD: RDD[_], finalNumPartitions: Int): Int = { var demand = finalNumPartitions @@ -1523,29 +1523,28 @@ private[spark] class DAGScheduler( } /** - * Up-front gang admission for an all-pipelined job (v1 / M1), checked BEFORE any stage exists. - * Because a v1 job is either all-regular or all-pipelined (mixed jobs are already rejected), an - * all-pipelined job's whole stage graph is one pipelined group with no regular prefix, so its - * full demand is known up front and the group is ready to admit immediately. Checking here - * (rather than in `submitStage` once a producer is already running) is true all-or-nothing gang - * admission: the whole group is admitted, or the job is failed before any member runs, so a - * member is never left running while a sibling cannot get slots -- and, like the barrier slot - * check, a rejection leaves no partial scheduler state. + * Up-front gang admission for an all-pipelined job, checked BEFORE any stage exists. Because a + * job is either all-regular or all-pipelined (mixed jobs are already rejected), an all-pipelined + * job's whole stage graph is one pipelined group with no regular prefix, so its full demand is + * known up front and the group is ready to admit immediately. Checking here (rather than in + * `submitStage` once a producer is already running) is true all-or-nothing gang admission: the + * whole group is admitted, or the job is failed before any member runs, so a member is never left + * running while a sibling cannot get slots -- and, like the barrier slot check, a rejection + * leaves no partial scheduler state. * - * Free-slot accounting follows spec S4.1: demand vs. total capacity (`maxNumConcurrentTasks` for - * the default profile -- v1 requires a single-profile group, spec S9) minus what OTHER work - * (other jobs) has outstanding -- running plus enqueued -- in that profile. Counting enqueued, - * not just running, tasks charges a busy neighbor's queued backlog against capacity too, so a - * group is not admitted against slots other work is already committed to. Fails the job via - * `listener` with no retry -- a transient shortfall is the caller's to retry (e.g. the streaming - * batch loop reruns the batch). Returns true (job failed) if it does not fit. + * Free-slot accounting: demand vs. total capacity (`maxNumConcurrentTasks` for the default + * profile -- a group is required to be single-profile) minus what OTHER work (other jobs) has + * outstanding -- running plus enqueued -- in that profile. Counting enqueued, not just running, + * tasks charges a busy neighbor's queued backlog against capacity too, so a group is not admitted + * against slots other work is already committed to. Fails the job via `listener` with no retry -- + * a transient shortfall is the caller's to retry (e.g. the streaming batch loop reruns it). + * Returns true (job failed) if it does not fit. * * The likeness to barrier's slot check is only that both reject before any stage is created * (leaving no partial state) and compute demand from the RDD graph. Retry behavior differs * deliberately: barrier RE-POSTS the job and re-runs its check on a timer up to * `spark.scheduler.barrier.maxConcurrentTasksCheck.maxFailures` times; pipelined admission is - * TERMINAL (one check, then fail), delegating transient-shortfall retry to the caller - * (scheduler-side PG-admission retry is a post-v1 refinement; spec S4.1). + * TERMINAL (one check, then fail), delegating transient-shortfall retry to the caller. * * The check can be turned off with `spark.scheduler.pipelinedGroup.slotCheck.enabled=false` (for * deployments that admit capacity out-of-band, e.g. via a slot reservation). @@ -1612,9 +1611,9 @@ private[spark] class DAGScheduler( * pipelined dependency. * * NOTE: this is deliberately defined here alongside the other group-topology helper - * (`isPipelinedProducer`). Its call sites are the group-atomic failure handling added later in - * this stack: tagging the member's `TaskSet.isPipelined` at submission and routing a member - * FetchFailed to a whole-group abort (spec S6). + * (`isPipelinedProducer`). Its call sites are the group-atomic failure handling: tagging the + * 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)) { @@ -1788,19 +1787,18 @@ private[spark] class DAGScheduler( return } - // v1 (M1, real-time-mode scope) supports a job that is either ALL-regular or ALL-pipelined, - // not a mix. A job whose shuffle graph combines a pipelined shuffle with a regular one is - // rejected up front (before any stage is created). Restricting to a single kind makes an - // all-pipelined job's whole stage graph one pipelined group with no regular prefix, so gang - // admission is decided up front (see the slot check below) rather than mid-DAG once a producer - // is already running. Mixed / mid-DAG groups are a later milestone. + // A job must be either ALL-regular or ALL-pipelined, not a mix: a job whose shuffle graph + // combines a pipelined shuffle with a regular one is rejected up front (before any stage is + // created). Restricting to a single kind makes an all-pipelined job's whole stage graph one + // pipelined group with no regular prefix, so gang admission is decided up front (see the slot + // check below) rather than mid-DAG once a producer is already running. val (hasPipelined, hasRegular) = classifyJobShuffleKinds(finalRDD) if (hasPipelined && hasRegular) { logWarning(log"Rejecting job ${MDC(JOB_ID, jobId)}: a job mixing a pipelined shuffle with " + log"a regular shuffle is not supported") listener.jobFailed(new SparkException( "A job that mixes a pipelined shuffle dependency with a regular shuffle dependency is " + - "not supported: v1 requires a job to be either all-regular or all-pipelined.")) + "not supported: a job must be either all-regular or all-pipelined.")) return } @@ -1892,7 +1890,7 @@ private[spark] class DAGScheduler( // shuffle to produce its map-output statistics. A PipelinedShuffleDependency cannot serve that: // it is transient with no durable, addressable map output, so it registers no map outputs and // getStatistics would return all-zero stats (misleading the map-stage/AQE consumer); and its - // producer/consumer co-scheduling makes speculation unsafe (spec S9). Reject a pipelined + // producer/consumer co-scheduling makes speculation unsafe. Reject a pipelined // dependency submitted as a map-stage job outright, up front, before any stage is created so no // partial scheduler state is left behind. Inert for a regular ShuffleDependency. (The // result-job path, handleJobSubmitted, only rejects the speculation case, since a pipelined @@ -2013,11 +2011,11 @@ 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 (v1 gang admission). + // unnecessary once admission is decided up front (gang admission). 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, - // so its successful completions are deferred until those producers finish (S5). + // so its successful completions are deferred until those producers finish. val deferral = dependentStageMap.getOrElseUpdate(stage, DependentStageInfo()) deferral.parents ++= pipelinedMissing @@ -2825,15 +2823,14 @@ private[spark] class DAGScheduler( val stage = stageIdToStage(task.stageId) - // Group-observable completion (spec S5), fine-grained model: a pipelined consumer's per-task - // side effects (accumulator update, task-end listener event) run in real time as its tasks - // finish -- exactly as for any other stage -- and ONLY the stage/job-completion decision is - // deferred until the producer(s) finish. Deferring the completion decision is what matters: - // advancing it early would let a consumer that finished ahead of its producer complete the job - // and cancel the still-running producer (via cancelRunningIndependentStages), or expose the - // consumer's output before the producer's. Its per-task facts, by contrast, are true when they - // happen (S5.1: task-level listener events flow in real time), so they must not be frozen for - // the producer's remaining lifetime. + // Group-observable completion for a pipelined consumer: its per-task side effects (accumulator + // update, task-end listener event) run in real time as its tasks finish -- exactly as for any + // other stage -- and ONLY the stage/job-completion decision is deferred until the producer(s) + // finish. Deferring the completion decision is what matters: advancing it early would let a + // consumer that finished ahead of its producer complete the job and cancel the still-running + // producer (via cancelRunningIndependentStages), or expose the consumer's output before the + // producer's. Its per-task facts, by contrast, are true when they happen (a task-end listener + // event flows in real time), so they must not be frozen for the producer's remaining lifetime. // // Two flags drive this. `finishOnly` is set on a replayed event (see // releaseDeferredPipelinedConsumers): its per-task effects already ran inline on first @@ -2879,7 +2876,7 @@ private[spark] class DAGScheduler( // Now that the per-task effects have run, defer the completion bookkeeping if this is a // co-scheduled pipelined consumer whose producer(s) are still running. It is replayed - // (finishOnly = true) once the last producer finishes, or dropped if a producer fails (S6). + // (finishOnly = true) once the last producer finishes, or dropped if a producer fails. if (deferCompletion) { val deferral = dependentStageMap(stage) logInfo(log"Deferring completion of task ${MDC(TASK_ID, event.taskInfo.taskId)} in " + @@ -3946,15 +3943,14 @@ private[spark] class DAGScheduler( * consumer was deferred on, remove it from that consumer's pending-producer set. When a consumer * has no pending producers left, either replay its buffered completion events (producer succeeded) * or drop them (a producer failed -- the group will be torn down and rerun, so the consumer's - * buffered completions must not be applied; S6). Inert unless `finishedStage` is a tracked - * producer. + * buffered completions must not be applied). Inert unless `finishedStage` is a tracked producer. * * Only the deferred stage/job-completion bookkeeping is buffered here; each event's per-task side * effects (accumulator update, task-end listener event) already ran inline when the task first - * completed (fine-grained model, S5). So replay re-posts each event with `finishOnly = true` (run - * the completion bookkeeping only, never the per-task effects again), and the drop path simply - * discards the buffered events -- there is nothing left to emit, since the consumer's TaskEnd - * events were already delivered in real time. + * completed. So replay re-posts each event with `finishOnly = true` (run the completion + * bookkeeping only, never the per-task effects again), and the drop path simply discards the + * buffered events -- there is nothing left to emit, since the consumer's TaskEnd events were + * already delivered in real time. */ private def releaseDeferredPipelinedConsumers( finishedStage: Stage, producerFailed: Boolean): Unit = { @@ -3978,12 +3974,11 @@ private[spark] class DAGScheduler( // Only the deferred stage/job-completion bookkeeping was buffered; the consumer's // per-task effects (TaskEnd, accumulator updates) already ran inline when the tasks first // completed. The group failed and will be rerun, so that completion bookkeeping must NOT - // be applied -- simply discard the buffered events. (There is nothing to re-emit: unlike - // the coarse model, no TaskEnd was withheld.) The already-applied inline accumulator - // deltas are NOT un-merged here -- consistent with base Spark, whose accumulators are - // non-transactional across an abort+rerun; the rerun re-delivers under the S5 idempotent- - // sink contract, and RTM SQL metrics are fresh per batch so there is no cross-batch - // double-count. + // be applied -- simply discard the buffered events. (There is nothing to re-emit: no + // TaskEnd was withheld.) The already-applied inline accumulator deltas are NOT un-merged + // here -- consistent with base Spark, whose accumulators are non-transactional across an + // abort+rerun; the rerun re-delivers to an idempotent sink, and streaming SQL metrics are + // fresh per batch so there is no cross-batch double-count. } else { logInfo(log"Replaying ${MDC(NUM_EVENTS, events.size.toLong)} deferred completion(s) for " + log"pipelined consumer ${MDC(STAGE, consumer)} now that its producers have finished") diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGSchedulerEvent.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGSchedulerEvent.scala index 58481e0517891..3c69fe5f34827 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGSchedulerEvent.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGSchedulerEvent.scala @@ -94,10 +94,10 @@ private[scheduler] case class CompletionEvent( metricPeaks: Array[Long], taskInfo: TaskInfo, // Set only on a pipelined consumer's completion event that was deferred and is now being - // replayed after its producer(s) finished (spec S5, fine-grained model). Its per-task side - // effects (accumulator update, task-end listener event) already ran inline when the task first - // completed; on replay only the deferred stage/job-completion bookkeeping must run, so this - // flag tells handleTaskCompletion to skip the per-task effects and avoid applying them twice. + // replayed after its producer(s) finished. Its per-task side effects (accumulator update, + // task-end listener event) already ran inline when the task first completed; on replay only the + // deferred stage/job-completion bookkeeping must run, so this flag tells handleTaskCompletion + // to skip the per-task effects and avoid applying them twice. finishOnly: Boolean = false) extends DAGSchedulerEvent 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 57ed13aa1798f..32607cf9dc03d 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -387,7 +387,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti @volatile var maxConcurrentTasksForTest: Int = 1000 override protected def maxConcurrentTasksForProfile(rpId: Int): Int = maxConcurrentTasksForTest - // Seam for the free-slot admission check (spec S4.1): outstanding (running + enqueued) task + // Seam for the free-slot admission check: outstanding (running + enqueued) task // demand of OTHER work. Default 0 so existing tests see a fully-free cluster; a test can set it // to model a busy neighbor. @volatile var outstandingTasksForOtherWorkForTest: (Int, Set[Int]) => Int = (_, _) => 0 @@ -6165,7 +6165,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } // ========================================================================================== - // Pipelined shuffle dependency: group formation + concurrent submission (M1.2) + // Pipelined shuffle dependency: group formation + concurrent submission // ========================================================================================== test("pipelined shuffle: consumer stage is submitted concurrently with its producer") { @@ -6216,7 +6216,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } test("pipelined shuffle: a job mixing a pipelined and a regular shuffle is rejected up front") { - // v1 (M1) supports a job that is either all-regular or all-pipelined, not a mix. A consumer + // A job must be either all-regular or all-pipelined, not a mix. A consumer // depending on BOTH a pipelined producer AND a regular producer is a mixed job and must be // rejected up front (before any stage is submitted), leaving no scheduler state behind. val pipelinedProducerRdd = new MyRDD(sc, 2, Nil) @@ -6398,7 +6398,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } // ========================================================================================== - // Gang admission / slot check (spec S4, S4.1) + // Gang admission / slot check // ========================================================================================== test("pipelined shuffle: an all-pipelined group that fits is admitted up front and runs") { @@ -6493,8 +6493,8 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } - test("pipelined shuffle: a group that fits total capacity but not FREE slots fails fast (S4.1)") { - // Spec S4.1: admission is decided against currently-FREE slots, not total capacity. A group of + test("pipelined shuffle: a group that fits total capacity but not FREE slots fails fast") { + // Admission is decided against currently-FREE slots, not total capacity. A group of // producer(2) + consumer(2) = 4 tasks fits a 10-slot cluster in principle, but if 8 slots are // already occupied by OTHER work only 2 are free -- the group cannot co-fit right now and must // fail fast rather than queue forever. Under the old total-capacity check this would wrongly be @@ -6516,7 +6516,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti submit(consumerRdd, Array(0, 1), listener = failListener) assert(failure.get() != null, - "a group that fits total but not free slots must fail the job (S4.1 free-slot check)") + "a group that fits total but not free slots must fail the job (free-slot check)") assert(failure.get().getMessage.contains("CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT") || failure.get().getMessage.contains("currently free"), s"expected a free-slot insufficient-slot error, got: ${failure.get().getMessage}") @@ -6581,7 +6581,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } // ========================================================================================== - // Group-observable (coarse deferred) completion (spec S5) + // Group-observable completion // ========================================================================================== test("pipelined shuffle: a consumer that finishes early does not end the job or cancel its " + @@ -6632,9 +6632,9 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti test("pipelined shuffle: deferred consumer completions are dropped when the producer fails") { // If a producer fails while a co-scheduled consumer's completion is deferred, the deferred // stage/job-completion bookkeeping must be DROPPED (not applied): the group is torn down / the - // job fails, and applying a partial consumer success would be incorrect (S6). The consumer's - // per-task side effects (TaskEnd) already ran inline when its tasks finished (fine-grained - // model, S5.1), so on the drop path there is nothing to re-emit and no TaskEnd is duplicated. + // job fails, and applying a partial consumer success would be incorrect. The consumer's + // per-task side effects (TaskEnd) already ran inline when its tasks finished, so on the drop + // path there is nothing to re-emit and no TaskEnd is duplicated. val taskEndCount = new java.util.concurrent.atomic.AtomicInteger(0) val countingListener = new SparkListener { override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = taskEndCount.incrementAndGet() @@ -6680,7 +6680,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti test("pipelined shuffle: a deferred consumer task fires its TaskEnd exactly once, in real time") { - // Fine-grained deferral (spec S5.1): a deferred consumer's per-task side effects (task-end + // A deferred consumer's per-task side effects (task-end // listener event, accumulator update) run in REAL TIME as its tasks finish -- NOT withheld // until replay -- and each runs exactly once (the completion bookkeeping alone is deferred, // then replayed with finishOnly=true, which must not re-post the TaskEnd). The producer (2 @@ -6701,7 +6701,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti val consumerTaskSet = taskSets(1) // Consumer finishes first. Its completion BOOKKEEPING is buffered (no result yet), but its - // per-task TaskEnd events fire immediately (fine-grained model). + // per-task TaskEnd events fire immediately. complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) assert(results.isEmpty, "consumer job result must be deferred until the producer finishes") sc.listenerBus.waitUntilEmpty(10000) From 43f50d0e55cc8b874220cacbbe7287d661878db3 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Thu, 23 Jul 2026 06:39:27 +0000 Subject: [PATCH 19/54] [SPARK-XXXXX][CORE] Reject a pipelined job whose member uses a non-default resource profile Up-front gang admission measures capacity and occupancy against the default resource profile, but each stage derives its profile from its RDDs. A pipelined job whose members use a non-default resource profile could therefore pass admission against the default profile's free slots and then queue or deadlock in the profile where its tasks actually run. Reject such a job up front in rejectUnadmittablePipelinedGroup (before any stage is created, like the other up-front rejections): walk the group's RDD graph and fail the job if any member carries a non-default resource profile. The whole pipelined group is required to run on the default profile -- the assumption the admission check already relies on. Per-profile admission is a follow-up. Add a DAGSchedulerSuite test: a pipelined consumer on a custom resource profile is rejected before any stage is created. Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 23 +++++++++++++ .../spark/scheduler/DAGSchedulerSuite.scala | 34 +++++++++++++++++++ 2 files changed, 57 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 34bbae09e35a3..5710c43c2f68d 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -1551,6 +1551,29 @@ private[spark] class DAGScheduler( */ private def rejectUnadmittablePipelinedGroup( jobId: Int, finalRDD: RDD[_], partitions: Array[Int], listener: JobListener): Boolean = { + // The whole group must run on the 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 + // therefore be admitted against the default profile's free slots yet run in a different, often + // smaller pool -- and could then queue or deadlock there. Reject such a job up front (before + // any stage is created, like the checks above), rather than admit it against the wrong pool. + var offendingRp: Option[ResourceProfile] = None + traverseRDDGraph(finalRDD) { (rdd, enqueue) => + 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 (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)) { return false } 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 32607cf9dc03d..ffde634534f75 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6339,6 +6339,40 @@ 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("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 061786f40eed6d311991deeb337defe6ea5ae06d Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Thu, 23 Jul 2026 16:23:52 +0000 Subject: [PATCH 20/54] [SPARK-XXXXX][CORE] Simplify pipelined-consumer deferral: defer the whole completion event, remove finishOnly Group-observable completion (spec S5) held a pipelined consumer that finished ahead of its still-running producer. The previous "fine-grained" model ran the consumer's per-task side effects (accumulator update, task-end listener event) inline and deferred only the stage/job-completion decision; making the side effects run exactly once then required a CompletionEvent.finishOnly flag so the replayed event would skip the already-applied per-task half, plus a special guard so a replay landing on a torn-down stage would not re-post TaskEnd. Switch to the simpler coarse model, per review feedback on #57341: defer the *whole* CompletionEvent and return before ANY of its side effects run. This makes the side effects run exactly once by construction -- at replay, when releaseDeferredPipelinedConsumers re-posts the buffered event and it re-enters handleTaskCompletion normally -- so: - CompletionEvent.finishOnly is removed entirely. - handleTaskCompletion has a single deferral check up front (before updateAccumulators / postTaskEnd), and the normal side-effect path is unguarded again (identical to base Spark), so the torn-down-stage postTaskEnd guard is no longer needed. - The producer-failure drop path emits the buffered tasks' TaskEnd events (they genuinely succeeded, so active-task-tracking listeners must see them finish) but does NOT run the stage/job-completion bookkeeping, since the group is rerun. Trade-off (accepted): a consumer that finishes early has its task-level listener events (TaskStart already fired; TaskEnd, accumulator updates) held until the producer's fate is known, rather than flowing in real time. The stage/job-completion decision -- the part that would wrongly cancel the still-running producer or expose output early -- is deferred either way; only the per-task event *timing* differs. Tests: the two deferral tests now assert the coarse timing (no consumer TaskEnd until the producer finishes; drop path emits TaskEnd but no result; replay still yields exactly 4 TaskEnd events with no buffer+replay duplication). Reverting the change (emitting TaskEnd inline before the deferral buffer) makes them fail (TaskEnd count 2 instead of 0 while the producer runs). Full DAGSchedulerSuite green (344 tests, both variants). Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 149 +++++++----------- .../spark/scheduler/DAGSchedulerEvent.scala | 8 +- .../spark/scheduler/DAGSchedulerSuite.scala | 54 +++---- 3 files changed, 87 insertions(+), 124 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 5710c43c2f68d..7a9f72d0589fe 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -2831,14 +2831,7 @@ private[spark] class DAGScheduler( // are properly notified and can chose to handle it. For instance, some listeners are // doing their own accounting and if they don't get the task end event they think // tasks are still running when they really aren't. - // - // Exception: a replayed pipelined-consumer completion (finishOnly) already fired its TaskEnd - // inline when the task first completed; only the deferred stage/job bookkeeping was - // re-posted. If the stage was torn down (e.g. a sibling aborted the group) before this replay - // dequeues, re-posting TaskEnd here would double-count it in listeners. Skip for finishOnly. - if (!event.finishOnly) { - postTaskEnd(event) - } + postTaskEnd(event) // Skip all the actions if the stage has been cancelled. return @@ -2846,68 +2839,59 @@ private[spark] class DAGScheduler( val stage = stageIdToStage(task.stageId) - // Group-observable completion for a pipelined consumer: its per-task side effects (accumulator - // update, task-end listener event) run in real time as its tasks finish -- exactly as for any - // other stage -- and ONLY the stage/job-completion decision is deferred until the producer(s) - // finish. Deferring the completion decision is what matters: advancing it early would let a - // consumer that finished ahead of its producer complete the job and cancel the still-running - // producer (via cancelRunningIndependentStages), or expose the consumer's output before the - // producer's. Its per-task facts, by contrast, are true when they happen (a task-end listener - // event flows in real time), so they must not be frozen for the producer's remaining lifetime. - // - // Two flags drive this. `finishOnly` is set on a replayed event (see - // releaseDeferredPipelinedConsumers): its per-task effects already ran inline on first - // completion, so on replay we skip straight to the completion bookkeeping and never apply them - // twice. `deferCompletion` is true for a fresh Success from a pipelined consumer whose - // producers are still running: run the per-task effects now, then buffer the event and return - // before the completion bookkeeping. Both are inert for a job with no pipelined dependency (the - // map is empty), so the regular path is unchanged. - val deferCompletion = !event.finishOnly && event.reason == Success && - dependentStageMap.get(stage).exists(_.parents.nonEmpty) - - if (!event.finishOnly) { - // Make sure the task's accumulators are updated before any other processing happens, so that - // we can post a task end event before any jobs or stages are updated. The accumulators are - // only updated in certain cases. - event.reason match { - case Success => - task match { - case rt: ResultTask[_, _] => - val resultStage = stage.asInstanceOf[ResultStage] - resultStage.activeJob match { - case Some(job) => - // Only update the accumulator once for each result task. - if (!job.finished(rt.outputId)) { - updateAccumulators(event) - } - case None => // Ignore update if task's job has finished. - } - case _ => - updateAccumulators(event) - } - case _: ExceptionFailure | _: TaskKilled => updateAccumulators(event) + // Group-observable completion (spec S5): if this stage is a pipelined consumer co-scheduled + // with a still-running pipelined producer, defer its *successful* completion in full until the + // producer(s) finish. Buffer the whole CompletionEvent and return before ANY of its side + // effects run (accumulator update, task-end listener event, stage/job completion) -- else a + // consumer finishing ahead of its producer would advance job completion and cancel the + // still-running producer (via cancelRunningIndependentStages), or expose its output early. + // Deferring the entire event (not just the completion bookkeeping) is what makes the side + // effects run exactly ONCE, at replay: releaseDeferredPipelinedConsumers re-posts the buffered + // event once the last producer completes (or drops it if a producer fails, S6), and it then + // re-enters here and runs the side effects normally. This deferral check must therefore precede + // updateAccumulators and postTaskEnd. Inert for jobs with no pipelined dependency (the map is + // empty), so the regular path is unchanged. + if (event.reason == Success) { + dependentStageMap.get(stage) match { + case Some(deferral) if deferral.parents.nonEmpty => + logInfo(log"Deferring completion of task ${MDC(TASK_ID, event.taskInfo.taskId)} in " + + log"pipelined consumer ${MDC(STAGE, stage)} until its producer(s) " + + log"${MDC(MISSING_PARENT_STAGES, deferral.parents.toSeq)} finish") + deferral.delayedTaskCompletionEvents += event + return case _ => } - if (trackingCacheVisibility) { - // Update rdd blocks' visibility status. - blockManagerMaster.updateRDDBlockVisibility( - event.taskInfo.taskId, visible = event.reason == Success) - } - - postTaskEnd(event) } - // Now that the per-task effects have run, defer the completion bookkeeping if this is a - // co-scheduled pipelined consumer whose producer(s) are still running. It is replayed - // (finishOnly = true) once the last producer finishes, or dropped if a producer fails. - if (deferCompletion) { - val deferral = dependentStageMap(stage) - logInfo(log"Deferring completion of task ${MDC(TASK_ID, event.taskInfo.taskId)} in " + - log"pipelined consumer ${MDC(STAGE, stage)} until its producer(s) " + - log"${MDC(MISSING_PARENT_STAGES, deferral.parents.toSeq)} finish") - deferral.delayedTaskCompletionEvents += event - return + // Make sure the task's accumulators are updated before any other processing happens, so that + // we can post a task end event before any jobs or stages are updated. The accumulators are + // only updated in certain cases. + event.reason match { + case Success => + task match { + case rt: ResultTask[_, _] => + val resultStage = stage.asInstanceOf[ResultStage] + resultStage.activeJob match { + case Some(job) => + // Only update the accumulator once for each result task. + if (!job.finished(rt.outputId)) { + updateAccumulators(event) + } + case None => // Ignore update if task's job has finished. + } + case _ => + updateAccumulators(event) + } + case _: ExceptionFailure | _: TaskKilled => updateAccumulators(event) + case _ => } + if (trackingCacheVisibility) { + // Update rdd blocks' visibility status. + blockManagerMaster.updateRDDBlockVisibility( + event.taskInfo.taskId, visible = event.reason == Success) + } + + postTaskEnd(event) event.reason match { case Success => @@ -3964,16 +3948,10 @@ private[spark] class DAGScheduler( /** * Called when a stage finishes. If `finishedStage` is a pipelined producer that some co-scheduled * consumer was deferred on, remove it from that consumer's pending-producer set. When a consumer - * has no pending producers left, either replay its buffered completion events (producer succeeded) - * or drop them (a producer failed -- the group will be torn down and rerun, so the consumer's - * buffered completions must not be applied). Inert unless `finishedStage` is a tracked producer. - * - * Only the deferred stage/job-completion bookkeeping is buffered here; each event's per-task side - * effects (accumulator update, task-end listener event) already ran inline when the task first - * completed. So replay re-posts each event with `finishOnly = true` (run the completion - * bookkeeping only, never the per-task effects again), and the drop path simply discards the - * buffered events -- there is nothing left to emit, since the consumer's TaskEnd events were - * already delivered in real time. + * has no pending producers left, either replay its buffered completion events (producer + * succeeded) or drop them (a producer failed -- the group will be torn down and rerun, so the + * consumer's buffered successes must not be applied; S6). Inert unless `finishedStage` is a + * tracked producer. */ private def releaseDeferredPipelinedConsumers( finishedStage: Stage, producerFailed: Boolean): Unit = { @@ -3994,20 +3972,15 @@ private[spark] class DAGScheduler( logInfo(log"Dropping ${MDC(NUM_EVENTS, events.size.toLong)} deferred completion(s) for " + log"pipelined consumer ${MDC(STAGE, consumer)} because its producer " + log"${MDC(STAGE, finishedStage)} failed; the group will be rerun") - // Only the deferred stage/job-completion bookkeeping was buffered; the consumer's - // per-task effects (TaskEnd, accumulator updates) already ran inline when the tasks first - // completed. The group failed and will be rerun, so that completion bookkeeping must NOT - // be applied -- simply discard the buffered events. (There is nothing to re-emit: no - // TaskEnd was withheld.) The already-applied inline accumulator deltas are NOT un-merged - // here -- consistent with base Spark, whose accumulators are non-transactional across an - // abort+rerun; the rerun re-delivers to an idempotent sink, and streaming SQL metrics are - // fresh per batch so there is no cross-batch double-count. + // The buffered tasks genuinely succeeded, so still emit their TaskEnd events -- otherwise + // listeners that track active tasks (e.g. AppStatusListener) would believe these tasks + // are still running. We deliberately do NOT run the stage/job completion bookkeeping: + // the group failed and will be rerun, so the consumer's results must not be applied. + events.foreach(postTaskEnd) } else { - logInfo(log"Replaying ${MDC(NUM_EVENTS, events.size.toLong)} deferred completion(s) for " + - log"pipelined consumer ${MDC(STAGE, consumer)} now that its producers have finished") - // Re-post each event to run ONLY the deferred completion bookkeeping (finishOnly = true); - // the per-task effects already ran inline and must not be applied again. - events.foreach(e => eventProcessLoop.post(e.copy(finishOnly = true))) + logInfo(log"Replaying ${MDC(NUM_EVENTS, events.size.toLong)} deferred completion(s) " + + log"for pipelined consumer ${MDC(STAGE, consumer)} now that its producers finished") + events.foreach(eventProcessLoop.post) } } } diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGSchedulerEvent.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGSchedulerEvent.scala index 3c69fe5f34827..cc788e5c65bcc 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGSchedulerEvent.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGSchedulerEvent.scala @@ -92,13 +92,7 @@ private[scheduler] case class CompletionEvent( result: Any, accumUpdates: Seq[AccumulatorV2[_, _]], metricPeaks: Array[Long], - taskInfo: TaskInfo, - // Set only on a pipelined consumer's completion event that was deferred and is now being - // replayed after its producer(s) finished. Its per-task side effects (accumulator update, - // task-end listener event) already ran inline when the task first completed; on replay only the - // deferred stage/job-completion bookkeeping must run, so this flag tells handleTaskCompletion - // to skip the per-task effects and avoid applying them twice. - finishOnly: Boolean = false) + taskInfo: TaskInfo) extends DAGSchedulerEvent private[scheduler] case class ExecutorAdded(execId: String, host: String) extends DAGSchedulerEvent 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 ffde634534f75..a3b6eb47ac4a2 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6664,11 +6664,12 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } test("pipelined shuffle: deferred consumer completions are dropped when the producer fails") { - // If a producer fails while a co-scheduled consumer's completion is deferred, the deferred - // stage/job-completion bookkeeping must be DROPPED (not applied): the group is torn down / the - // job fails, and applying a partial consumer success would be incorrect. The consumer's - // per-task side effects (TaskEnd) already ran inline when its tasks finished, so on the drop - // path there is nothing to re-emit and no TaskEnd is duplicated. + // If a producer fails while a co-scheduled consumer's completions are buffered, those buffered + // successes must be DROPPED (not applied): the group is torn down / the job fails, and applying + // a partial consumer success would be incorrect (S6). The buffered tasks genuinely succeeded, + // so their TaskEnd events are still emitted on the drop path (otherwise listeners that track + // active tasks would believe these tasks are still running); only the stage/job-completion + // bookkeeping is withheld. val taskEndCount = new java.util.concurrent.atomic.AtomicInteger(0) val countingListener = new SparkListener { override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = taskEndCount.incrementAndGet() @@ -6687,25 +6688,25 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti val producerTaskSet = taskSets.head val consumerTaskSet = taskSets(1) - // Consumer finishes early -> its completion bookkeeping is deferred, but its 2 TaskEnd events - // fire in real time. + // Consumer finishes early -> its whole completion events are buffered (deferred): no result + // and no TaskEnd yet, since the entire event is withheld until the producer's fate is known. complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) - assert(results.isEmpty, "consumer job result should be deferred while the producer runs") + assert(results.isEmpty, "consumer completions should be buffered while the producer runs") sc.listenerBus.waitUntilEmpty(10000) - assert(taskEndCount.get() === 2, - "consumer's TaskEnd events must fire inline, before the producer's fate; got " + + assert(taskEndCount.get() === 0, + "buffered consumer completions must not fire TaskEnd until released; got " + taskEndCount.get()) - // The producer now FAILS its whole task set. The job fails; the deferred consumer completion - // must NOT be applied as a result, and no consumer TaskEnd is re-emitted on the drop path. + // The producer now FAILS its whole task set. The job fails; the buffered consumer successes + // must NOT be applied as results, but their TaskEnd events ARE emitted on the drop path so + // active-task-tracking listeners see the tasks finish. failed(producerTaskSet, "producer blew up") assert(failure.get() != null, "job should fail when the producer fails") assert(results.isEmpty, - "deferred consumer completion must be dropped when the producer fails, not applied") + "buffered consumer successes must be dropped when the producer fails, not applied") sc.listenerBus.waitUntilEmpty(10000) assert(taskEndCount.get() === 2, - "drop path must not re-emit consumer TaskEnd events (they already fired); got " + - taskEndCount.get()) + "drop path must emit the buffered consumer TaskEnd events; got " + taskEndCount.get()) assertDataStructuresEmpty() } finally { sc.removeSparkListener(countingListener) @@ -6713,14 +6714,11 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } - test("pipelined shuffle: a deferred consumer task fires its TaskEnd exactly once, in real time") { - // A deferred consumer's per-task side effects (task-end - // listener event, accumulator update) run in REAL TIME as its tasks finish -- NOT withheld - // until replay -- and each runs exactly once (the completion bookkeeping alone is deferred, - // then replayed with finishOnly=true, which must not re-post the TaskEnd). The producer (2 - // tasks) and consumer (2 tasks) yield exactly 4 TaskEnd events total; a buffer+replay - // double-post would inflate that to 6, and a withhold-until-replay bug would delay the 2 - // consumer events. + 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 + // update) applied exactly once -- at replay -- not once when buffered and again when replayed. + // The producer (2 tasks) and consumer (2 tasks) yield exactly 4 TaskEnd events total; a + // buffer+replay double-post would inflate that count. val taskEndCount = new java.util.concurrent.atomic.AtomicInteger(0) val countingListener = new SparkListener { override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = taskEndCount.incrementAndGet() @@ -6734,17 +6732,15 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti val producerStageId = taskSets.head.stageId val consumerTaskSet = taskSets(1) - // Consumer finishes first. Its completion BOOKKEEPING is buffered (no result yet), but its - // per-task TaskEnd events fire immediately. + // Consumer finishes first -> its two completions are buffered (deferred, no TaskEnd yet). complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) assert(results.isEmpty, "consumer job result must be deferred until the producer finishes") sc.listenerBus.waitUntilEmpty(10000) - assert(taskEndCount.get() === 2, - "consumer's 2 TaskEnd events must fire in real time, not at replay; got " + + assert(taskEndCount.get() === 0, + "deferred consumer completions must not fire TaskEnd until replay; got " + taskEndCount.get()) - // Producer finishes -> the two deferred consumer completions replay (finishOnly=true): the - // job result is applied now, but NO additional TaskEnd is posted (they already fired above). + // Producer finishes -> the two deferred consumer completions replay. completeShuffleMapStageSuccessfully(producerStageId, 0, 2) assert(results === Map(0 -> 42, 1 -> 43)) sc.listenerBus.waitUntilEmpty(10000) From 148c5b35ab666c1438b49d0581dfd5faed60f25e Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Thu, 23 Jul 2026 17:36:31 +0000 Subject: [PATCH 21/54] [SPARK-XXXXX][CORE] Flush a buffered consumer's TaskEnds on job teardown, not just on the release path The coarse deferral model buffers a pipelined consumer's whole completion event -- including its TaskEnd -- until group completion. releaseDeferredPipelinedConsumers flushes those buffered TaskEnds on both its replay (producer succeeded) and drop (producer failed) paths, and it is normally reached for every producer before job cleanup runs (cancelRunningIndependentStages finishes each running producer first). But cleanupStateForJobAndIndependentStages -> removeStage discarded a consumer's deferral with a bare `dependentStageMap -= stage`, without flushing its buffered TaskEnds. That silently loses them in a case the release path does not cover: a producer left in resubmit limbo -- finished with no error yet not available, so producerAboutToResubmit retained the deferral AND markStageAsFinished removed the producer from runningStages -- is skipped by cancelRunningIndependentStages (it is neither running nor failed). If the job is then cancelled or aborted, cleanup drops the consumer's buffered completion without ever emitting its TaskEnds, so a listener that tracks active tasks (e.g. AppStatusListener) leaks them as perpetually running. This was impossible before the coarse-model change, when TaskEnd fired inline on first completion and nothing was buffered to lose. Fix: at removeStage, drain the removed consumer's buffered delayedTaskCompletionEvents through postTaskEnd (as the producer-failed drop path does) before discarding the deferral. The buffered successes are still not applied as results -- this is a job failure / cancellation -- only their TaskEnds are flushed. Test: a regression test drives a producer into resubmit limbo, cancels the job, and asserts the 2 buffered consumer TaskEnds are flushed (and no result applied). Reverting the fix makes it fail (0 TaskEnds instead of 2). Full DAGSchedulerSuite green. Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 14 ++++- .../spark/scheduler/DAGSchedulerSuite.scala | 54 +++++++++++++++++++ 2 files changed, 67 insertions(+), 1 deletion(-) 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 7a9f72d0589fe..e54216ac98b46 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -1134,7 +1134,19 @@ private[spark] class DAGScheduler( // Drop any pipelined-completion deferral keyed on this stage (as consumer), and // remove it from other consumers' pending-producer sets (as producer), so no // deferral outlives its job (e.g. on job abort before the producers finished). - dependentStageMap -= stage + // A consumer's buffered completion events hold TaskEnds that have not been emitted + // yet (the whole event is deferred until group completion). This teardown is a job + // failure / cancellation, so the buffered successes must NOT be applied as results + // -- but their TaskEnds must still be flushed, exactly as the producer-failed drop + // path does (see releaseDeferredPipelinedConsumers), or a listener tracking active + // tasks (e.g. AppStatusListener) would leak these tasks as perpetually running. + // Normally releaseDeferredPipelinedConsumers has already drained the deferral by + // the time cleanup runs (cancelRunningIndependentStages finishes each running + // producer first); this covers the case where it has not -- e.g. a producer left in + // resubmit limbo (finished but not available) that cancelRunningIndependentStages + // skips because it is neither running nor failed. + dependentStageMap.remove(stage).foreach(_.delayedTaskCompletionEvents.foreach( + postTaskEnd)) dependentStageMap.values.foreach(_.parents -= stage) } // data structures based on StageId 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 a3b6eb47ac4a2..330949ae6969a 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6803,6 +6803,60 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } + test("pipelined shuffle: job teardown flushes a buffered consumer's TaskEnds even when the " + + "release path never drained it") { + // A consumer's whole completion event is buffered while its producer runs (coarse model), so + // its TaskEnds are held, not yet emitted. Normally releaseDeferredPipelinedConsumers drains + // them (cancelRunningIndependentStages finishes each running producer before cleanup). But a + // producer left in resubmit limbo -- finished with no error yet NOT available, so + // producerAboutToResubmit held the deferral AND markStageAsFinished removed it from + // runningStages -- is skipped by cancelRunningIndependentStages (neither running nor failed). + // If the job is then cancelled, cleanup must still flush the buffered TaskEnds (results stay + // dropped) so a listener tracking active tasks does not leak them as perpetually running. + 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 jobId = submit(consumerRdd, Array(0, 1)) + val producerStageId = taskSets.head.stageId + val consumerTaskSet = taskSets(1) + + // Consumer finishes early -> whole completion buffered; no TaskEnd yet. + complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) + sc.listenerBus.waitUntilEmpty() + assert(taskEndCount.get() === 0, "buffered consumer TaskEnds are held while the producer runs") + + // Drive the producer into resubmit limbo: finished, no error, but not available -> the + // deferral is retained and the producer leaves runningStages. No task failed, so it also has + // no failedAttemptIds -- cancelRunningIndependentStages will skip it. + val producerStage = scheduler.stageIdToStage(producerStageId).asInstanceOf[ShuffleMapStage] + assert(!producerStage.isAvailable, "precondition: producer not available") + scheduler.markStageAsFinished(producerStage, errorMessage = None, willRetry = false) + assert(scheduler.dependentStageMap.keys.exists(_.rdd eq consumerRdd), + "precondition: the deferral is retained (producer about to resubmit)") + assert(!scheduler.runningStages.contains(producerStage) && + producerStage.failedAttemptIds.isEmpty, + "precondition: producer is neither running nor failed, so cancel will skip it") + + // Cancel the job. The release path never drains this consumer, so cleanup must flush it. + cancel(jobId) + sc.listenerBus.waitUntilEmpty() + assert(scheduler.dependentStageMap.isEmpty, "the deferral must not outlive the job") + assert(results.isEmpty, "a cancelled job's buffered consumer success must not be applied") + assert(taskEndCount.get() === 2, + s"teardown must flush the 2 buffered consumer TaskEnds (else active tasks leak); got " + + taskEndCount.get()) + assertDataStructuresEmpty() + } finally { + sc.removeSparkListener(countingListener) + } + } + } class DAGSchedulerAbortStageOffSuite extends DAGSchedulerSuite { From 97f415285db7d3bdb4e3999e9a2c4f8322ee89ae Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Thu, 23 Jul 2026 22:06:23 +0000 Subject: [PATCH 22/54] [SPARK-58263][CORE] Address review nits: correct two scheduler comments Two documentation-only corrections from review (cloud-fan), no behavior change: - TaskSchedulerImpl.outstandingTasksForOtherWorkInProfile scaladoc: "the slots free by everything else" reversed the relationship (other work consumes slots, it does not free them). Reword to "the slots left free after accounting for everything else in the SAME resource profile". - DAGScheduler.handleMapStageSubmitted comment: the parenthetical said the result-job path (handleJobSubmitted) "only rejects the speculation case". It rejects both speculation and dynamic allocation via rejectUnsupportedPipelinedJob; correct the summary. Co-authored-by: Isaac --- .../scala/org/apache/spark/scheduler/DAGScheduler.scala | 5 +++-- .../org/apache/spark/scheduler/TaskSchedulerImpl.scala | 6 +++--- 2 files changed, 6 insertions(+), 5 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 e54216ac98b46..1cba1b8855ebd 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -1928,8 +1928,9 @@ private[spark] class DAGScheduler( // producer/consumer co-scheduling makes speculation unsafe. Reject a pipelined // dependency submitted as a map-stage job outright, up front, before any stage is created so no // partial scheduler state is left behind. Inert for a regular ShuffleDependency. (The - // result-job path, handleJobSubmitted, only rejects the speculation case, since a pipelined - // dependency is a legitimate internal edge there.) + // result-job path, handleJobSubmitted, rejects the speculation and dynamic-allocation cases via + // rejectUnsupportedPipelinedJob, but a pipelined dependency is otherwise a legitimate internal + // edge there.) if (dependency.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) { logWarning(log"Rejecting map-stage job ${MDC(JOB_ID, jobId)}: a pipelined shuffle dependency " + log"cannot be materialized as a map-stage job") diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala index b832bc6463f16..196ae04d8acd8 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala @@ -165,9 +165,9 @@ private[spark] class TaskSchedulerImpl( * the stages in `excludeStageIds`. "Outstanding" is the not-yet-completed demand of each task set * -- running plus enqueued (`numTasks - tasksSuccessful`) -- not just the tasks actively running, * so a neighbor's queued backlog is charged against capacity too. Used by the pipelined-group - * slot admission check (see DAGScheduler): it compares the group's demand against the slots free - * by everything else in the SAME resource profile, so it excludes the group's own members (whose - * tasks would otherwise be charged against the group's own admission). + * slot admission check (see DAGScheduler): it compares the group's demand against the slots left + * free after accounting for everything else in the SAME resource profile, so it excludes the + * group's own members (whose tasks would otherwise be charged against the group's own admission). * * Computed from the TaskSetManagers so it is resource-profile-scoped, skips zombie (superseded) * attempts so a retried stage is not double-counted, and is taken under a single lock so the From 23a3be15689b5663c75246c2093638e1d3683253 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Fri, 24 Jul 2026 00:37:18 +0000 Subject: [PATCH 23/54] [SPARK-XXXXX][CORE] Fix scalastyle line-length in the teardown-flush test assert The assert message in "job teardown flushes a buffered consumer's TaskEnds" was 101 chars, over the 100-char scalastyle limit for test sources. Trim it. 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 330949ae6969a..30e0f24e2efbd 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6829,7 +6829,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti // Consumer finishes early -> whole completion buffered; no TaskEnd yet. complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) sc.listenerBus.waitUntilEmpty() - assert(taskEndCount.get() === 0, "buffered consumer TaskEnds are held while the producer runs") + assert(taskEndCount.get() === 0, "buffered consumer TaskEnds are held while producer runs") // Drive the producer into resubmit limbo: finished, no error, but not available -> the // deferral is retained and the producer leaves runningStages. No task failed, so it also has From 25167363d11650198ba308cd71e0ae16979e7ac4 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Fri, 24 Jul 2026 17:56:53 +0000 Subject: [PATCH 24/54] [SPARK-58263][CORE] Clarify that buffering the whole event defers all per-task side effects Address review (cloud-fan): make the group-observable-completion comment explicit that buffering the whole CompletionEvent DEFERS all three side effects -- the accumulator update, the task-end listener event, and stage/job completion -- not just the stage/job-completion bookkeeping. Also state the listener-visible consequence: a consumer task that finishes ahead of its producer emits no SparkListenerTaskEnd until the producer finishes, so the Spark UI reports that already-succeeded task as still running for the producer's remaining lifetime. Comment-only; no behavior change. Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 22 +++++++++++-------- 1 file changed, 13 insertions(+), 9 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 1cba1b8855ebd..a4308752fc2b1 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -2855,15 +2855,19 @@ private[spark] class DAGScheduler( // Group-observable completion (spec S5): if this stage is a pipelined consumer co-scheduled // with a still-running pipelined producer, defer its *successful* completion in full until the // producer(s) finish. Buffer the whole CompletionEvent and return before ANY of its side - // effects run (accumulator update, task-end listener event, stage/job completion) -- else a - // consumer finishing ahead of its producer would advance job completion and cancel the - // still-running producer (via cancelRunningIndependentStages), or expose its output early. - // Deferring the entire event (not just the completion bookkeeping) is what makes the side - // effects run exactly ONCE, at replay: releaseDeferredPipelinedConsumers re-posts the buffered - // event once the last producer completes (or drops it if a producer fails, S6), and it then - // re-enters here and runs the side effects normally. This deferral check must therefore precede - // updateAccumulators and postTaskEnd. Inert for jobs with no pipelined dependency (the map is - // empty), so the regular path is unchanged. + // effects run. Buffering the whole event DEFERS all three of those effects -- the accumulator + // update, the task-end listener event, and stage/job completion -- until the event is replayed; + // it is NOT just the stage/job-completion bookkeeping that waits. (A listener-visible + // consequence: a consumer task that finishes ahead of its producer emits no SparkListenerTaskEnd + // until the producer finishes, so the Spark UI reports that already-succeeded task as still + // running for the producer's remaining lifetime.) Else a consumer finishing ahead of its + // producer would advance job completion and cancel the still-running producer (via + // cancelRunningIndependentStages), or expose its output early. + // Deferring the entire event is what makes the side effects run exactly ONCE, at replay: + // releaseDeferredPipelinedConsumers re-posts the buffered event once the last producer completes + // (or drops it if a producer fails, S6), and it then re-enters here and runs the side effects + // normally. This deferral check must therefore precede updateAccumulators and postTaskEnd. Inert + // for jobs with no pipelined dependency (the map is empty), so the regular path is unchanged. if (event.reason == Success) { dependentStageMap.get(stage) match { case Some(deferral) if deferral.parents.nonEmpty => From eebb3298a37b4c78aa94cce57291dea148fd955b Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sat, 25 Jul 2026 01:52:32 +0000 Subject: [PATCH 25/54] [SPARK-58263][CORE] Drop a fan-in consumer's deferral as soon as any producer fails releaseDeferredPipelinedConsumers decided drop-vs-replay only when a consumer's last pending producer finished, keyed on that producer's outcome. For a consumer co-scheduled with more than one pipelined producer (fan-in), producer A could fail while B was still pending -- A was removed from the pending set but no drop occurred, and A's failure was forgotten -- so when B later succeeded it emptied the set on the replay branch and the consumer's buffered successes were replayed even though they were computed against A's now-invalid output. Fix: when a producer fails, drop the consumer's deferral immediately -- flush its buffered TaskEnds and remove the entry -- regardless of whether it still has other pending producers. A pipelined group is failed as a unit (S6), so once any producer fails the buffered successes must not be applied. Dropping now (rather than remembering the failure until the last producer finishes) also leaves no deferral state that could outlive the call and be reused by a later re-co-scheduling of the same stage. A surviving producer that finishes afterward simply finds the consumer already gone (a clean no-op). Producer success is unchanged: remove the finished producer and replay only once the last one has succeeded. Recovery from a drop, if any, is a caller rerun of the whole group (a new job, S6); the scheduler does not rerun the group in place. Test: a two-producer (fan-in) case where A fails and B then succeeds, asserting the consumer's deferral is dropped as soon as A fails and B's later success does not replay it. Reverting to the old parents-empty gating makes it fail (the successes replay). Full DAGSchedulerSuite green. Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 48 +++++++++++-------- .../spark/scheduler/DAGSchedulerSuite.scala | 44 +++++++++++++++++ 2 files changed, 72 insertions(+), 20 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 a4308752fc2b1..5288b74a6950f 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -3964,11 +3964,16 @@ private[spark] class DAGScheduler( /** * Called when a stage finishes. If `finishedStage` is a pipelined producer that some co-scheduled - * consumer was deferred on, remove it from that consumer's pending-producer set. When a consumer - * has no pending producers left, either replay its buffered completion events (producer - * succeeded) or drop them (a producer failed -- the group will be torn down and rerun, so the - * consumer's buffered successes must not be applied; S6). Inert unless `finishedStage` is a - * tracked producer. + * consumer was deferred on, resolve that consumer's deferral: + * - Producer FAILED: drop the consumer's buffered completions immediately, regardless of whether + * it still has other pending producers. A pipelined group is failed as a unit (S6), so once a + * producer fails the consumer's buffered successes must not be applied -- and dropping now, + * rather than remembering the failure until the last producer finishes, means no failure state + * outlives this call (a surviving producer that later finishes just finds the consumer gone). + * - Producer SUCCEEDED: remove it from the consumer's pending-producer set, and replay the + * buffered completions only once the LAST producer has succeeded (the set becomes empty). + * Inert unless `finishedStage` is a tracked producer. Recovery from a drop, if any, is a caller + * rerun of the whole group (a new job, S6) -- the scheduler does not rerun the group in place. */ private def releaseDeferredPipelinedConsumers( finishedStage: Stage, producerFailed: Boolean): Unit = { @@ -3982,23 +3987,26 @@ private[spark] class DAGScheduler( for (consumer <- released) { val deferral = dependentStageMap(consumer) deferral.parents -= finishedStage - if (deferral.parents.isEmpty) { + if (producerFailed) { + // Drop immediately: the group is being failed as a unit, so this consumer's buffered + // successes must not be applied no matter what its other producers do. Removing the entry + // now leaves no deferral state to go stale for a later re-co-scheduling of this stage. dependentStageMap -= consumer val events = deferral.delayedTaskCompletionEvents.toList - if (producerFailed) { - logInfo(log"Dropping ${MDC(NUM_EVENTS, events.size.toLong)} deferred completion(s) for " + - log"pipelined consumer ${MDC(STAGE, consumer)} because its producer " + - log"${MDC(STAGE, finishedStage)} failed; the group will be rerun") - // The buffered tasks genuinely succeeded, so still emit their TaskEnd events -- otherwise - // listeners that track active tasks (e.g. AppStatusListener) would believe these tasks - // are still running. We deliberately do NOT run the stage/job completion bookkeeping: - // the group failed and will be rerun, so the consumer's results must not be applied. - events.foreach(postTaskEnd) - } else { - logInfo(log"Replaying ${MDC(NUM_EVENTS, events.size.toLong)} deferred completion(s) " + - log"for pipelined consumer ${MDC(STAGE, consumer)} now that its producers finished") - events.foreach(eventProcessLoop.post) - } + logInfo(log"Dropping ${MDC(NUM_EVENTS, events.size.toLong)} deferred completion(s) for " + + log"pipelined consumer ${MDC(STAGE, consumer)} because a producer failed; failing the " + + log"group") + // The buffered tasks genuinely succeeded, so still emit their TaskEnd events -- otherwise + // listeners that track active tasks (e.g. AppStatusListener) would believe these tasks are + // still running. We deliberately do NOT run the stage/job completion bookkeeping: the group + // is being failed as a unit, so the consumer's results must not be applied. + events.foreach(postTaskEnd) + } else if (deferral.parents.isEmpty) { + dependentStageMap -= consumer + val events = deferral.delayedTaskCompletionEvents.toList + logInfo(log"Replaying ${MDC(NUM_EVENTS, events.size.toLong)} deferred completion(s) " + + log"for pipelined consumer ${MDC(STAGE, consumer)} now that its producers finished") + events.foreach(eventProcessLoop.post) } } } 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 30e0f24e2efbd..8fb8d27c83981 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6713,6 +6713,50 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } } + test("pipelined shuffle: a fan-in consumer's buffered successes are dropped as soon as one " + + "producer fails, even if another later succeeds") { + // Fan-in: a consumer co-scheduled with TWO pipelined producers (A and B). If A fails, the + // consumer's buffered successes must be DROPPED -- they were computed against A's now-invalid + // output, and the group is failed as a unit (S6). The drop happens immediately when A fails, + // not deferred until the last producer finishes: so even if B later succeeds, the consumer has + // already been dropped and B's completion is a clean no-op (it must NOT replay the successes). + val producerA = new MyRDD(sc, 2, Nil) + val psdA = new PipelinedShuffleDependency(producerA, new HashPartitioner(2)) + val producerB = new MyRDD(sc, 2, Nil) + val psdB = new PipelinedShuffleDependency(producerB, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(psdA, psdB), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + assert(taskSets.size === 3, s"A, B and the consumer must be co-scheduled; got ${taskSets.size}") + val consumerTaskSet = + taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd).get + val stageA = scheduler.stageIdToStage( + taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq producerA).get.stageId) + val stageB = scheduler.stageIdToStage( + taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq producerB).get.stageId) + val consumerStage = scheduler.stageIdToStage(consumerTaskSet.stageId) + + // Consumer finishes early -> its successes are buffered against BOTH A and B. + complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) + assert(results.isEmpty, "consumer completions should be buffered while its producers run") + assert(scheduler.dependentStageMap.get(consumerStage).exists(_.parents.size == 2), + "the consumer must be deferred against both producers") + + // A fails: the consumer is dropped immediately (not retained pending B), and no result applied. + // (Driven via markStageAsFinished directly; the job-failure teardown that would normally follow + // is exercised by the group-atomic-failure tests in a later PR -- here we assert only the + // drop-vs-replay outcome.) + scheduler.markStageAsFinished(stageA, errorMessage = Some("producer A blew up")) + assert(!scheduler.dependentStageMap.contains(consumerStage), + "the consumer's deferral must be dropped as soon as a producer fails") + assert(results.isEmpty, "a fan-in consumer's buffered successes must not be applied on failure") + + // B now succeeds: the consumer is already gone, so this is a clean no-op -- it must NOT replay. + scheduler.runningStages += stageB + completeShuffleMapStageSuccessfully(stageB.id, 0, 2) + assert(results.isEmpty, + s"a dropped fan-in consumer must not replay when a surviving producer succeeds; got $results") + } + 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 From 787869ba4263a977fa8e87f2a80ee1e0e5e08469 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sat, 25 Jul 2026 04:28:13 +0000 Subject: [PATCH 26/54] [SPARK-58263][CORE] Document the barrier-resubmit edge for the fan-in drop path The immediate-drop on producer failure treats producerFailed as terminal. That holds for every producer-failure path that reaches releaseDeferredPipelinedConsumers except one: a barrier stage that fails a task and is then resubmitted (its handler calls markStageAsFinished(errorMessage) without willRetry=true, then resubmits). For a fan-in consumer, dropping its deferral for a barrier producer that will retry loses the consumer's buffered results. A barrier stage in a pipelined group is rejected up front by a later PR in this stack (checkPipelinedProducerSupported), which removes this path entirely; on this PR standalone it is a known, narrow limitation (needs a barrier producer feeding a fan-in consumer). Document it inline rather than adding the rejection early. Comment-only; no behavior change. Co-authored-by: Isaac --- .../scala/org/apache/spark/scheduler/DAGScheduler.scala | 7 +++++++ 1 file changed, 7 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 5288b74a6950f..c4bf3786aca64 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -3991,6 +3991,13 @@ private[spark] class DAGScheduler( // Drop immediately: the group is being failed as a unit, so this consumer's buffered // successes must not be applied no matter what its other producers do. Removing the entry // now leaves no deferral state to go stale for a later re-co-scheduling of this stage. + // This treats producerFailed as terminal, which holds for every producer-failure path that + // reaches here EXCEPT a barrier stage that fails a task and is resubmitted (its handler + // calls markStageAsFinished(errorMessage) without willRetry=true, then resubmits): + // there, dropping a fan-in consumer's deferral for a producer that will retry loses that + // consumer's buffered results. A barrier stage in a pipelined group is rejected up front by + // a later PR in this stack (checkPipelinedProducerSupported), which removes this path; on + // this PR standalone it is a known, narrow limitation (barrier producer + fan-in consumer). dependentStageMap -= consumer val events = deferral.delayedTaskCompletionEvents.toList logInfo(log"Dropping ${MDC(NUM_EVENTS, events.size.toLong)} deferred completion(s) for " + From 4c842ec88b884de0d9ed75cd61eb459093f538ad Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sat, 11 Jul 2026 05:33:58 +0000 Subject: [PATCH 27/54] [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 c4bf3786aca64..245bda85ff7f8 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -2447,7 +2447,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 78e4fc78b8801..9c483ad67406c 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala @@ -69,6 +69,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 @@ -1105,16 +1114,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 } } @@ -1190,7 +1214,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 8fb8d27c83981..d6e8f8a607c4a 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6239,6 +6239,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) @@ -6901,6 +6942,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 9006cde0a4d6347dd8b2c59931d9aa7f7c840569 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sat, 11 Jul 2026 19:20:49 +0000 Subject: [PATCH 28/54] [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 b17f8c5facad3..cc433684c4680 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -6319,6 +6319,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 245bda85ff7f8..b0dc8634e7e0d 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 => @@ -2988,7 +3003,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 d6e8f8a607c4a..848054591f722 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6974,6 +6974,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 24aeed9840795597c6c52c7b383737263ca100d5 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sat, 11 Jul 2026 19:55:57 +0000 Subject: [PATCH 29/54] [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 b0dc8634e7e0d..2907cfdafe972 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -3095,6 +3095,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) @@ -3211,39 +3236,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 => @@ -3783,6 +3776,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 848054591f722..5567117dcd184 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7263,6 +7263,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 dd481984237a4cf1f99e4f4a63fecc85f871bda3 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 13 Jul 2026 04:27:04 +0000 Subject: [PATCH 30/54] [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 5567117dcd184..2fec30d1b3c53 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6241,7 +6241,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 @@ -6263,13 +6264,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 @@ -6950,8 +6953,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") @@ -6979,7 +6984,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. @@ -6998,7 +7004,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, @@ -7014,8 +7021,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)) @@ -7072,7 +7081,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". @@ -7082,7 +7092,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, @@ -7099,10 +7110,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) @@ -7118,7 +7131,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))) @@ -7141,7 +7155,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) @@ -7164,7 +7179,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)") @@ -7184,10 +7200,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)) @@ -7268,7 +7286,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 @@ -7314,7 +7333,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. @@ -7333,7 +7353,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 @@ -7357,7 +7378,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 @@ -7376,7 +7398,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), @@ -7389,6 +7412,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 f4b6449a6fbf7a810dd46bac1cab2b672a391c89 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sun, 19 Jul 2026 06:20:27 +0000 Subject: [PATCH 31/54] [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 2fec30d1b3c53..bac4c3023f85b 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6239,47 +6239,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() } @@ -7329,90 +7310,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 7b94634c1c6b16767ae081957f192a0715f148a9 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sun, 19 Jul 2026 19:17:58 +0000 Subject: [PATCH 32/54] [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 2907cfdafe972..d6a31dc3570a1 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -3116,7 +3116,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 bac4c3023f85b..fe803394568f7 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7147,15 +7147,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 102dae5f3fdff8b1bd9401785242cf651d73d1ea Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 06:29:46 +0000 Subject: [PATCH 33/54] [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 d6a31dc3570a1..6381ae01198fe 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -3007,7 +3007,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 @@ -3023,8 +3024,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) { @@ -3099,7 +3100,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 9c483ad67406c..3237a73185645 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala @@ -1216,12 +1216,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 fe803394568f7..4d8a888c95c55 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6962,7 +6962,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" + @@ -7038,15 +7039,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)) @@ -7078,7 +7079,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 => @@ -7090,7 +7091,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 @@ -7098,7 +7099,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) @@ -7123,7 +7124,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 @@ -7134,8 +7135,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 @@ -7168,7 +7169,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 => @@ -7266,8 +7267,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 @@ -7305,7 +7306,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() @@ -7314,7 +7315,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 f82342371ec1a8c4524b9b42102395e07be41389 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 06:54:19 +0000 Subject: [PATCH 34/54] [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 6381ae01198fe..c035137b28784 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -1904,6 +1904,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)}) " + @@ -2460,9 +2463,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 8ef0c963768e10ab4016da563ceafba567050d87 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 17:35:34 +0000 Subject: [PATCH 35/54] [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 1de7a5dedeeb66d7b28ff5a5fa280a263c0ade28 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 22:18:10 +0000 Subject: [PATCH 36/54] [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 4d8a888c95c55..411934cc5da7e 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7313,6 +7313,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 88adca5be72a39df11c3118c2ac9265da8f81763 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 22:33:55 +0000 Subject: [PATCH 37/54] [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 411934cc5da7e..6a7da2426dbc9 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6291,6 +6291,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 4197dece18b42319bdc3b82e7c8ea021ab670f6f Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 23:09:44 +0000 Subject: [PATCH 38/54] [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 6a7da2426dbc9..53eb65a723915 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7342,6 +7342,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 efa26432a12fe498f192d141be5a8b9a749bf77b Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Tue, 21 Jul 2026 07:17:10 +0000 Subject: [PATCH 39/54] [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 c035137b28784..1a7c57d323d72 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)) { @@ -3012,7 +3012,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 @@ -3028,7 +3028,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. @@ -3104,7 +3104,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 @@ -3124,7 +3124,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 53eb65a723915..11e786be28080 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6241,10 +6241,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) @@ -6991,7 +6991,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 // ========================================================================================== @@ -7211,15 +7211,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)) @@ -7269,7 +7268,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)) @@ -7297,11 +7296,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. @@ -7440,16 +7438,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. @@ -7481,10 +7478,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 8bdf00a51ccae74aa6db9256736d4e6952535327 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Thu, 23 Jul 2026 17:15:55 +0000 Subject: [PATCH 40/54] [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 11e786be28080..4c6910eefdeef 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7371,72 +7371,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 183d9bbf246e1692e2d0926a54ad5bf8553e7dff Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 13 Jul 2026 04:17:58 +0000 Subject: [PATCH 41/54] [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 cc433684c4680..4104d04df7660 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -6325,6 +6325,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 1a7c57d323d72..6b5fa9e27c360 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.{ResourceProfile, TaskResourceProfile} import org.apache.spark.resource.ResourceProfile.{DEFAULT_RESOURCE_PROFILE_ID, EXECUTOR_CORES_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 @@ -2065,7 +2167,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 eb3c1da539571439d89238b44d3d71ca6636fb48 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 13 Jul 2026 04:45:36 +0000 Subject: [PATCH 42/54] [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 6b5fa9e27c360..24ea56d4b6164 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 */ @@ -1962,9 +2011,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) @@ -1997,6 +2050,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 4c6910eefdeef..08a6ecc48df81 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 @@ -7443,6 +7443,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 e33cf33d1fa30e74bd6917d0d704b5fe2805032a Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sun, 19 Jul 2026 07:26:41 +0000 Subject: [PATCH 43/54] [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 08a6ecc48df81..8aee41d8cd2ee 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6811,6 +6811,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 @@ -7489,8 +7528,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 @@ -7588,7 +7627,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[_]])) @@ -7611,7 +7650,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 2514e363ffef3db7cc72f08f48eb167d6365ddee Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sun, 19 Jul 2026 19:19:07 +0000 Subject: [PATCH 44/54] [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 24ea56d4b6164..02d6c440089c1 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 16bbc1467743a3b8dc7a790d3dcb6d6057e3104a Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 06:16:23 +0000 Subject: [PATCH 45/54] [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 8aee41d8cd2ee..920d6e0e4cabe 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7483,18 +7483,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 = @@ -7519,25 +7517,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"), @@ -7545,11 +7541,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 19a5afc434875993612f778d0f5f215ebf12689a Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 06:39:35 +0000 Subject: [PATCH 46/54] [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 02d6c440089c1..e9b74d18d1f34 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 @@ -2053,10 +2064,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 @@ -4612,6 +4624,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 920d6e0e4cabe..a900f86104de6 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7649,6 +7649,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 a37a6266e5ff58c5d033d67e503526f8243b6501 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 17:34:41 +0000 Subject: [PATCH 47/54] [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 e9b74d18d1f34..a83a65b6a7b4c 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 @@ -3283,8 +3293,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 @@ -4627,8 +4639,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 a900f86104de6..5c427b4a0de39 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6813,13 +6813,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) @@ -7649,23 +7649,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 8388681228bc857388fe80195c0a666ccf26fb10 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 20 Jul 2026 21:21:47 +0000 Subject: [PATCH 48/54] [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 a83a65b6a7b4c..3efb630553090 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 @@ -1777,6 +1782,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 @@ -1840,24 +1855,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 @@ -2095,6 +2098,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)}) " + @@ -2249,6 +2257,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 0722df2edf3721f904475c9ea6f629c4066e988f Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Tue, 21 Jul 2026 01:33:06 +0000 Subject: [PATCH 49/54] [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 3efb630553090..f8b329d53b770 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -2256,17 +2256,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 c181dabaa8cfe8abcef3fce85f7736642909e6d8 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Tue, 21 Jul 2026 07:17:27 +0000 Subject: [PATCH 50/54] [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 f8b329d53b770..2f5afc1d4c9c7 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: @@ -2040,7 +2041,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. @@ -2100,8 +2101,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( @@ -2256,12 +2257,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 @@ -2270,16 +2271,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, @@ -3320,7 +3321,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 @@ -4661,9 +4662,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 5c427b4a0de39..74707ed8e87e4 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6814,8 +6814,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 @@ -6831,7 +6831,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 @@ -7485,14 +7485,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 = @@ -7561,7 +7561,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) @@ -7590,9 +7590,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) @@ -7610,7 +7610,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 @@ -7632,9 +7632,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) @@ -7649,9 +7649,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 c7870b33c7e454f6a50c7a7baf2dcd114868a0d8 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Thu, 23 Jul 2026 06:42:54 +0000 Subject: [PATCH 51/54] [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 | 23 ------------- .../spark/scheduler/DAGSchedulerSuite.scala | 34 ------------------- 2 files changed, 57 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 2f5afc1d4c9c7..c1d6f4888e451 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -1759,29 +1759,6 @@ private[spark] class DAGScheduler( */ private def rejectUnadmittablePipelinedGroup( jobId: Int, finalRDD: RDD[_], partitions: Array[Int], listener: JobListener): Boolean = { - // The whole group must run on the 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 - // therefore be admitted against the default profile's free slots yet run in a different, often - // smaller pool -- and could then queue or deadlock there. Reject such a job up front (before - // any stage is created, like the checks above), rather than admit it against the wrong pool. - var offendingRp: Option[ResourceProfile] = None - traverseRDDGraph(finalRDD) { (rdd, enqueue) => - 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 (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 74707ed8e87e4..2a648e94d8991 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -6393,40 +6393,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("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 0460e80ddebc465a482d72ce43895975ef0e0211 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Thu, 23 Jul 2026 17:22:57 +0000 Subject: [PATCH 52/54] [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 2a648e94d8991..085fa00c30a90 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7449,16 +7449,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 = @@ -7483,23 +7483,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"), @@ -7510,6 +7510,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 996fc1cae7f5967239d1b54929f8f3c1fc9eec2e Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Thu, 23 Jul 2026 19:18:03 +0000 Subject: [PATCH 53/54] [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 c1d6f4888e451..ee820b1b96639 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 @@ -2018,11 +2020,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 085fa00c30a90..8f93f59b0fd8e 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7687,6 +7687,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 2bd265710880f308a7e31d3d7735f7303842911f Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Fri, 24 Jul 2026 00:36:40 +0000 Subject: [PATCH 54/54] [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 8f93f59b0fd8e..3474478dba618 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7687,7 +7687,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