You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by we...@apache.org on 2026/09/20 12:22:18 UTC
(spark) branch branch-4.x updated: [SPARK-59670][CORE] Avoid redundant waiting-stage scans and job-preflight traversals
This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new e0785b310d28 [SPARK-59670][CORE] Avoid redundant waiting-stage scans and job-preflight traversals
e0785b310d28 is described below
commit e0785b310d28924d31e42100fd507828a59cde3e
Author: Boyang Jerry Peng <je...@databricks.com>
AuthorDate: Sun Sep 20 12:21:49 2026 +0000
[SPARK-59670][CORE] Avoid redundant waiting-stage scans and job-preflight traversals
### What changes were proposed in this pull request?
Reduce redundant DAGScheduler work introduced by pipelined-shuffle scheduling:
1. Return from `submitWaitingPipelinedChildStages` before scanning waiting stages when the
running stage is not a pipelined producer. Pipelined-producer wake-up behavior is unchanged.
2. Reuse the negative pipelined-dependency check when speculation or dynamic allocation is
enabled, instead of traversing an ordinary job's RDD graph again for shuffle-shape
classification. Preserve short-circuit rejection and validation precedence.
3. Count pipelined producer task demand and enqueue dependencies in one dependency loop,
preserving shuffle-ID deduplication and graph traversal order.
The changes are submission-local: no cross-job cache, new configuration, or scheduling policy.
JIRA: [SPARK-59670](https://issues.apache.org/jira/browse/SPARK-59670)
### Why are the changes needed?
The pipelined-child helper currently scans waiting stages even when the running stage cannot
have any pipelined children to wake. Repeated scans add avoidable event-loop work, especially
for large ordinary stage graphs.
Ordinary jobs with speculation or dynamic allocation enabled also pay for two negative
preflight graph traversals. The first has already established that no pipelined dependency
exists, so the second cannot add classification information. Pipelined task-demand calculation
separately enumerates each RDD's dependencies twice when one loop suffices.
This removes those redundant operations, not all scheduler graph traversals. With speculation
and dynamic allocation both disabled, ordinary-job preflight still uses its existing classifier.
### Does this PR introduce _any_ user-facing change?
No. Scheduling, admission decisions, error messages, and validation precedence are unchanged.
The changes reduce scheduler overhead.
### How was this patch tested?
Added 17 regression cases to `DAGSchedulerSuite` covering:
- No waiting-stage parent inspection when starting ordinary result or shuffle-map stages.
- A single ordinary-job preflight traversal across speculation/dynamic-allocation and
shuffle/no-shuffle combinations.
- Short-circuit feature rejection and precedence over mixed-shape rejection.
- Single dependency enumeration during pipelined task-demand calculation, with unchanged demand.
- Successful job completion and scheduler-state cleanup on the relevant paths.
Local validation with Java 17:
```bash
build/sbt -batch \
'core/testOnly org.apache.spark.scheduler.DAGSchedulerSuite' \
'core/Compile/scalastyle' \
'core/Test/scalastyle'
```
- All 256 `DAGSchedulerSuite` tests passed.
- Both production and test scalastyle checks passed with zero errors or warnings.
- Mutation checks temporarily restored the redundant operations and confirmed that the
corresponding traversal-count tests fail. The optimized code was restored and the full
suite rerun successfully.
- `git diff --check` passed.
An exploratory local helper benchmark used 1,000- and 10,000-RDD chains with a shuffle every
16 edges. Graph setup was outside timing. Separate counters verified that ordinary preflight
(with speculation enabled) and pipelined task-demand calculation each reduced dependency
enumeration from `2 * (N - 1)` to `N - 1`, with identical classification/demand results.
Best warmed ordinary-preflight timings were 59.7 to 25.7 microseconds for 1,000 RDDs and
1.096 to 0.513 milliseconds for 10,000 RDDs. These are helper-level results from one local run,
not end-to-end query speedups. No waiting-stage-scan timing benchmark was run, and the temporary
benchmark harness is not included in this PR.
### Was this patch authored or co-authored using generative AI tooling?
Co-authored with OpenAI Codex CLI 0.144.1
Closes #58931 from jerrypeng/SPARK-59670-dagscheduler-traversals.
Authored-by: Boyang Jerry Peng <je...@databricks.com>
Signed-off-by: Wenchen Fan <we...@databricks.com>
(cherry picked from commit 00d52e23e81088d78c3e46159270aee189de7d61)
Signed-off-by: Wenchen Fan <we...@databricks.com>
---
.../org/apache/spark/scheduler/DAGScheduler.scala | 69 ++++----
.../apache/spark/scheduler/DAGSchedulerSuite.scala | 176 ++++++++++++++++++++-
2 files changed, 215 insertions(+), 30 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 6d4d9c7a5b2f..73a20be14164 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
@@ -1164,8 +1164,8 @@ private[spark] class DAGScheduler(
}
/**
- * Shape of the shuffle boundaries in the RDD graph rooted at a job's final RDD, computed by
- * [[classifyJobShuffleShape]] before any stage is created (so an unsupported job is rejected
+ * Shape of the shuffle boundaries in the RDD graph rooted at a job's final RDD, computed
+ * before any stage is created (so an unsupported job is rejected
* fail-fast with no partial scheduler state).
*
* The supported shapes are:
@@ -1315,13 +1315,9 @@ private[spark] class DAGScheduler(
}
/**
- * 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.
+ * Classify a job's shuffle shape, or fail it via `listener` if it uses a pipelined shuffle with
+ * an incompatible cluster feature. Checked before any stage is created, so a rejection leaves
+ * no partial scheduler state. The map-stage-job path rejects pipelined dependencies outright.
*
* Rejected combinations:
* - Speculation: a speculative copy of a producer would race a consumer already reading the
@@ -1332,12 +1328,13 @@ private[spark] class DAGScheduler(
* 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.
+ *
+ * @return the shuffle shape, or None if the job was rejected for an incompatible feature
*/
- private def rejectUnsupportedPipelinedJob(
+ private def classifyJobShuffleShapeAndCheckFeatures(
jobId: Int,
finalRDD: RDD[_],
- listener: JobListener): Boolean = {
- // Only walk the RDD graph if a relevant feature is on, then only once.
+ listener: JobListener): Option[JobShuffleShape] = {
val speculationOn = sc.conf.get(config.SPECULATION_ENABLED)
val dynAllocOn = Utils.isDynamicAllocationEnabled(sc.conf)
if ((speculationOn || dynAllocOn) && rddGraphHasPipelinedDependency(finalRDD)) {
@@ -1347,7 +1344,7 @@ private[spark] class DAGScheduler(
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
+ return None
}
// dynAllocOn
logWarning(log"Rejecting job ${MDC(JOB_ID, jobId)}: dynamic allocation is incompatible " +
@@ -1356,9 +1353,18 @@ private[spark] class DAGScheduler(
"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
+ return None
+ }
+ if (speculationOn || dynAllocOn) {
+ // The negative feature check already proved there is no pipelined dependency anywhere.
+ // Reuse that result instead of walking the graph again to classify its shuffle shape.
+ Some(JobShuffleShape(
+ hasPipelined = false,
+ hasUnmaterializedRegularBoundary = false,
+ hasPipelinedBelowRegular = false))
+ } else {
+ Some(classifyJobShuffleShape(finalRDD))
}
- false
}
/**
@@ -1936,8 +1942,11 @@ private[spark] class DAGScheduler(
* stage.
*/
private def submitWaitingPipelinedChildStages(runningParent: Stage): Unit = {
+ if (!isPipelinedProducer(runningParent)) {
+ return
+ }
val pipelinedChildren = waitingStages.filter { child =>
- child.parents.contains(runningParent) && isPipelinedProducer(runningParent)
+ child.parents.contains(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).
@@ -2127,15 +2136,17 @@ private[spark] class DAGScheduler(
var demand = finalNumPartitions
val countedShuffleIds = new HashSet[Int]
traverseRDDGraph(finalRDD) { (rdd, enqueue) =>
- rdd.dependencies.foreach {
- case pd: PipelinedShuffleDependency[_, _, _] =>
- if (countedShuffleIds.add(pd.shuffleId)) {
- demand += pd.rdd.partitions.length
- }
- case _ => // narrow deps are not boundaries; a regular dep is a materialized prefix
- // boundary (JobShuffleShape), whose stages never run and are not members
+ rdd.dependencies.foreach { dep =>
+ dep match {
+ case pd: PipelinedShuffleDependency[_, _, _] =>
+ if (countedShuffleIds.add(pd.shuffleId)) {
+ demand += pd.rdd.partitions.length
+ }
+ case _ => // narrow deps are not boundaries; a regular dep is a materialized prefix
+ // boundary (JobShuffleShape), whose stages never run and are not members
+ }
+ enqueue(dep.rdd)
}
- rdd.dependencies.foreach(dep => enqueue(dep.rdd))
}
demand
}
@@ -2431,8 +2442,9 @@ private[spark] class DAGScheduler(
// 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
+ val shape = classifyJobShuffleShapeAndCheckFeatures(jobId, finalRDD, listener) match {
+ case Some(shape) => shape
+ case None => return
}
// A job's shuffle graph must be all-regular, all-pipelined, or the materialized-prefix mixed
@@ -2440,7 +2452,6 @@ private[spark] class DAGScheduler(
// JobShuffleShape). Anything else is rejected up front (before any stage is created): an
// unmaterialized regular stage would have to run while gang-admitted producers hold slots
// blocked on transport backpressure, which admission does not account for and can deadlock.
- val shape = classifyJobShuffleShape(finalRDD)
val hasPipelined = shape.hasPipelined
if (shape.isUnsupportedMix) {
logWarning(log"Rejecting job ${MDC(JOB_ID, jobId)}: a job mixing a pipelined shuffle with " +
@@ -2594,8 +2605,8 @@ private[spark] class DAGScheduler(
// 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, rejects the speculation and dynamic-allocation cases via
- // rejectUnsupportedPipelinedJob, but a pipelined dependency is otherwise a legitimate internal
- // edge there.)
+ // classifyJobShuffleShapeAndCheckFeatures, 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/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala
index 1d175c3d6ae6..7f17803f9a62 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala
@@ -6759,6 +6759,37 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti
// Pipelined shuffle dependency: group formation + concurrent submission
// ==========================================================================================
+ for (withShuffle <- Seq(false, true)) {
+ test(s"non-pipelined stage start skips waiting stages (withShuffle=$withShuffle)") {
+ val producerRdd = new MyRDD(sc, 2, Nil)
+ val finalRdd = if (withShuffle) {
+ val dep = new ShuffleDependency(producerRdd, new HashPartitioner(2))
+ new MyRDD(sc, 2, List(dep), tracker = mapOutputTracker)
+ } else {
+ producerRdd
+ }
+ val unrelatedWaitingStage = mock(classOf[Stage])
+ when(unrelatedWaitingStage.id).thenReturn(Int.MaxValue)
+ when(unrelatedWaitingStage.parents).thenReturn(Nil)
+ scheduler.waitingStages += unrelatedWaitingStage
+ try {
+ submit(finalRdd, Array(0, 1))
+ verify(unrelatedWaitingStage, never()).parents
+ assert(scheduler.waitingStages.contains(unrelatedWaitingStage))
+ } finally {
+ scheduler.waitingStages -= unrelatedWaitingStage
+ }
+
+ if (withShuffle) {
+ completeShuffleMapStageSuccessfully(taskSets(0).stageId, 0, 2)
+ }
+ val resultTaskSet = taskSets(if (withShuffle) 1 else 0)
+ completeAndCheckAnswer(
+ resultTaskSet, Seq((Success, 42), (Success, 43)), Map(0 -> 42, 1 -> 43))
+ assertDataStructuresEmpty()
+ }
+ }
+
test("pipelined shuffle: consumer stage is submitted concurrently with its producer") {
// producer (shuffle map) --[pipelined]--> consumer (result)
val producerRdd = new MyRDD(sc, 2, Nil)
@@ -6975,7 +7006,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti
}
test("pipelined shuffle: a non-mixed job is classified by the cheap kinds pre-pass") {
- // classifyJobShuffleShape runs on EVERY job submission, including on deployments that never
+ // Shuffle-shape preflight runs on every job submission, including on deployments that never
// enable this feature, so a non-mixed job must not pay the precise (RDD, belowRegular)-keyed
// walk that only a pipelined/regular MIX needs.
//
@@ -7107,6 +7138,115 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti
assertDataStructuresEmpty()
}
+ for {
+ speculation <- Seq(false, true)
+ dynamicAllocation <- Seq(false, true)
+ withShuffle <- Seq(false, true)
+ } {
+ test(s"ordinary job preflight traverses once (speculation=$speculation, " +
+ s"dynamicAllocation=$dynamicAllocation, withShuffle=$withShuffle)") {
+ val producerRdd = new MyRDD(sc, 2, Nil)
+ val inputRdd = if (withShuffle) {
+ val dep = new ShuffleDependency(producerRdd, new HashPartitioner(2))
+ new MyRDD(sc, 2, List(dep), tracker = mapOutputTracker)
+ } else {
+ producerRdd
+ }
+ val edgeReads = new AtomicInteger
+ val dep = new OneToOneDependency(inputRdd) {
+ override def rdd: RDD[(Int, Int)] = {
+ edgeReads.incrementAndGet()
+ super.rdd
+ }
+ }
+ val finalRdd = new MyRDD(sc, 2, List(dep))
+ var preflightEdgeReads = -1
+ doAnswer { invocation =>
+ // Snapshot before stage discovery adds its own dependency traversals.
+ if (preflightEdgeReads == -1) {
+ preflightEdgeReads = edgeReads.get()
+ }
+ invocation.callRealMethod()
+ }.when(scheduler).getShuffleDependenciesAndResourceProfiles(finalRdd)
+
+ sc.conf.set(config.SPECULATION_ENABLED, speculation)
+ sc.conf.set(config.DYN_ALLOCATION_ENABLED, dynamicAllocation)
+ sc.conf.set(config.DYN_ALLOCATION_TESTING, true)
+ try {
+ edgeReads.set(0)
+ submit(finalRdd, Array(0, 1))
+ assert(preflightEdgeReads === 1)
+ if (withShuffle) {
+ completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2)
+ }
+ completeAndCheckAnswer(
+ taskSets.last, Seq((Success, 42), (Success, 43)), Map(0 -> 42, 1 -> 43))
+ assertDataStructuresEmpty()
+ } finally {
+ sc.conf.set(config.SPECULATION_ENABLED, false)
+ sc.conf.set(config.DYN_ALLOCATION_ENABLED, false)
+ sc.conf.set(config.DYN_ALLOCATION_TESTING, false)
+ }
+ }
+ }
+
+ for ((speculation, dynamicAllocation) <- Seq((true, false), (false, true), (true, true))) {
+ val expectedError = if (speculation) {
+ "Speculative execution is not supported"
+ } else {
+ "Dynamic allocation is not supported"
+ }
+
+ test(s"pipelined shuffle: feature rejection short-circuits (speculation=$speculation, " +
+ s"dynamicAllocation=$dynamicAllocation)") {
+ val producerRdd = new MyRDD(sc, 2, Nil)
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+ val untouchedRdd = new MyRDD(sc, 2, Nil) {
+ override protected def getDependencies: Seq[Dependency[_]] =
+ throw new IllegalStateException("this branch must not be visited")
+ }
+ val consumerRdd = new MyRDD(sc, 2,
+ List(pipelinedDep, new OneToOneDependency(untouchedRdd)))
+ sc.conf.set(config.SPECULATION_ENABLED, speculation)
+ sc.conf.set(config.DYN_ALLOCATION_ENABLED, dynamicAllocation)
+ sc.conf.set(config.DYN_ALLOCATION_TESTING, true)
+ try {
+ val failure = submitAndCaptureFailure(consumerRdd, Array(0, 1))
+ assert(failure != null)
+ assert(failure.getMessage.contains(expectedError))
+ assert(taskSets.isEmpty)
+ assertDataStructuresEmpty()
+ } finally {
+ sc.conf.set(config.SPECULATION_ENABLED, false)
+ sc.conf.set(config.DYN_ALLOCATION_ENABLED, false)
+ sc.conf.set(config.DYN_ALLOCATION_TESTING, false)
+ }
+ }
+
+ test(s"pipelined shuffle: feature rejection precedes mixed-shape rejection " +
+ s"(speculation=$speculation, dynamicAllocation=$dynamicAllocation)") {
+ val producerRdd = new MyRDD(sc, 2, Nil)
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+ val middleRdd = new MyRDD(sc, 2, List(pipelinedDep))
+ val regularDep = new ShuffleDependency(middleRdd, new HashPartitioner(2))
+ val finalRdd = new MyRDD(sc, 2, List(regularDep))
+ sc.conf.set(config.SPECULATION_ENABLED, speculation)
+ sc.conf.set(config.DYN_ALLOCATION_ENABLED, dynamicAllocation)
+ sc.conf.set(config.DYN_ALLOCATION_TESTING, true)
+ try {
+ val failure = submitAndCaptureFailure(finalRdd, Array(0, 1))
+ assert(failure != null)
+ assert(failure.getMessage.contains(expectedError))
+ assert(taskSets.isEmpty)
+ assertDataStructuresEmpty()
+ } finally {
+ sc.conf.set(config.SPECULATION_ENABLED, false)
+ sc.conf.set(config.DYN_ALLOCATION_ENABLED, false)
+ sc.conf.set(config.DYN_ALLOCATION_TESTING, false)
+ }
+ }
+ }
+
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.
@@ -7242,6 +7382,40 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti
// Gang admission / slot check
// ==========================================================================================
+ test("pipelined shuffle: admission enumerates each dependency once") {
+ val producerRdd = new MyRDD(sc, 2, Nil)
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+ val deps = List(pipelinedDep)
+ val edgeVisits = new AtomicInteger
+ val countingDeps = new scala.collection.immutable.AbstractSeq[Dependency[_]] {
+ override def length: Int = deps.length
+ override def apply(index: Int): Dependency[_] = deps(index)
+ override def iterator: Iterator[Dependency[_]] = deps.iterator.map { dep =>
+ edgeVisits.incrementAndGet()
+ dep
+ }
+ }
+ val consumerRdd = new MyRDD(sc, 2, Nil) {
+ override protected def getDependencies: Seq[Dependency[_]] = countingDeps
+ }
+ val myScheduler = scheduler.asInstanceOf[MyDAGScheduler]
+ myScheduler.maxConcurrentTasksForTest = 0
+ try {
+ val failure = submitAndCaptureFailure(consumerRdd, Array(0, 1))
+ assert(failure.isInstanceOf[SparkException])
+ checkError(
+ exception = failure.asInstanceOf[SparkException],
+ condition = "CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT",
+ parameters = scala.collection.immutable.Map("numTasks" -> "4", "numSlots" -> "0"))
+ // Classification and admission each enumerate the edge once; no stage was created.
+ assert(edgeVisits.get() === 2)
+ assert(taskSets.isEmpty)
+ assertDataStructuresEmpty()
+ } finally {
+ myScheduler.maxConcurrentTasksForTest = 1000
+ }
+ }
+
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,
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org