dongjoon-hyun commented on code in PR #59523:
URL: https://github.com/apache/spark/pull/59523#discussion_r4349530202
##########
core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala:
##########
@@ -1516,6 +1515,29 @@ private[spark] class DAGScheduler(
checkPipelinedStageCreation(finalRDD, partitions, jobId)
}
+ /** Reject persistence that could bypass a native pipelined shuffle reader.
*/
+ private def checkPipelinedCacheSupported(finalRDD: RDD[_]): Unit = {
+ val persisted = new HashSet[RDD[_]]
+ traverseRDDGraph(finalRDD) { (rdd, enqueue) =>
+ if (rdd.getStorageLevel != StorageLevel.NONE ||
+
rdd.checkpointData.exists(_.isInstanceOf[LocalRDDCheckpointData[_]])) {
+ persisted += rdd
+ }
+ rdd.dependencies.foreach(dep => enqueue(dep.rdd))
+ }
+ // Cache locations cannot establish a stable cut: eviction can revive an
upstream producer,
+ // and a partial cache hit skips a reader whose producer still needs its
termination ACK.
+ // Check storage intent instead, including pending local checkpoints and
narrow descendants.
+ val visited = new HashSet[RDD[_]]
+ if (persisted.exists(rdd => rddChainReadsPipelinedShuffle(rdd, visited))) {
Review Comment:
**Spark SQL persists a pipelined consumer by itself, so this rejection now
breaks pandas-on-Spark's default index, with advice the user can't follow.**
With cache enabled, which is the default for the `distributed-sequence` index
type ([config.py
L177-L196](https://github.com/apache/spark/blob/39ebf2923484b045ef4b4f864db444431503ccd0/python/pyspark/pandas/config.py#L177-L196)),
`AttachDistributedSequenceExec` persists its child's RDD, or local-checkpoints
it when `compute.default_index_cache=LOCAL_CHECKPOINT`
([AttachDistributedSequenceExec.scala
L77-L96](https://github.com/apache/spark/blob/39ebf2923484b045ef4b4f864db444431503ccd0/sql/core/src/main/scala/org/apache/spark/sql/execution/python/AttachDistributedSequenceExec.scala#L77-L96)),
and then calls `zipWithIndex()` on it
([L100-L107](https://github.com/apache/spark/blob/39ebf2923484b045ef4b4f864db444431503ccd0/sql/core/src/main/scala/org/apache/spark/sql/execution/python/AttachDistributedSequenceExec.scala#L100-L107)).
`Zippe
dWithIndexRDD` stops at the persisted ancestor and runs a counting job on it
([ZippedWithIndexRDD.scala
L51](https://github.com/apache/spark/blob/39ebf2923484b045ef4b4f864db444431503ccd0/core/src/main/scala/org/apache/spark/rdd/ZippedWithIndexRDD.scala#L51),
[L76-L79](https://github.com/apache/spark/blob/39ebf2923484b045ef4b4f864db444431503ccd0/core/src/main/scala/org/apache/spark/rdd/ZippedWithIndexRDD.scala#L76-L79)).
`PipelinedShuffleEligibility` has this operator neither in
`isUnsupportedConsumer` nor in `hasUnsupportedBoundary`
([L49-L52](https://github.com/apache/spark/blob/39ebf2923484b045ef4b4f864db444431503ccd0/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/PipelinedShuffleEligibility.scala#L49-L52),
[L105-L110](https://github.com/apache/spark/blob/39ebf2923484b045ef4b4f864db444431503ccd0/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/PipelinedShuffleEligibility.scala#L105-L110)),
so the child exchange becomes pipelined, and the counting jo
b's final RDD is exactly what this check rejects.
For example, with the settings of `test_pipelined_shuffle.py`
([L28-L40](https://github.com/apache/spark/blob/39ebf2923484b045ef4b4f864db444431503ccd0/python/pyspark/sql/tests/test_pipelined_shuffle.py#L28-L40)),
`spark.range(0, 4000, 1, 2).repartition(4).pandas_api().to_pandas()` should
now fail with `PIPELINED_SHUFFLE_UNSUPPORTED`, with AQE on or off. I traced
this statically and didn't run it. Before this PR, the `LOCAL_CHECKPOINT` mode
worked because `doCheckpoint` cut the lineage, and the default mode worked
until a channel queue filled up, at which point the partial cache hit this PR
describes made it hang. The message asks to "unpersist cached RDDs", but this
RDD is created inside Spark, so the only workarounds are
`compute.default_index_cache=NONE` or turning the flag off.
`CacheManager` already keeps cache construction on regular shuffles for the
same reason ([CacheManager.scala
L701-L704](https://github.com/apache/spark/blob/39ebf2923484b045ef4b4f864db444431503ccd0/sql/core/src/main/scala/org/apache/spark/sql/execution/CacheManager.scala#L701-L704)).
Could `AttachDistributedSequenceExec` with `cache = true` be treated as an
unsupported boundary in `PipelinedShuffleEligibility`, in this PR or under a
separate JIRA, so that this plan stays regular?
##########
core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala:
##########
@@ -1516,6 +1515,29 @@ private[spark] class DAGScheduler(
checkPipelinedStageCreation(finalRDD, partitions, jobId)
}
+ /** Reject persistence that could bypass a native pipelined shuffle reader.
*/
+ private def checkPipelinedCacheSupported(finalRDD: RDD[_]): Unit = {
+ val persisted = new HashSet[RDD[_]]
+ traverseRDDGraph(finalRDD) { (rdd, enqueue) =>
+ if (rdd.getStorageLevel != StorageLevel.NONE ||
+
rdd.checkpointData.exists(_.isInstanceOf[LocalRDDCheckpointData[_]])) {
+ persisted += rdd
+ }
+ rdd.dependencies.foreach(dep => enqueue(dep.rdd))
+ }
+ // Cache locations cannot establish a stable cut: eviction can revive an
upstream producer,
+ // and a partial cache hit skips a reader whose producer still needs its
termination ACK.
+ // Check storage intent instead, including pending local checkpoints and
narrow descendants.
+ val visited = new HashSet[RDD[_]]
+ if (persisted.exists(rdd => rddChainReadsPipelinedShuffle(rdd, visited))) {
+ throw pipelinedUnsupportedError(
+ "cache/persist/localCheckpoint in a pipelined consumer's within-stage
chain " +
Review Comment:
**The error doesn't say which RDD has to change.** `persisted.exists(...)`
drops the RDD that matched, and the warning at L2756 only logs the job id. When
the same job also caches a producer input, which this PR explicitly allows, the
user has to work out by hand which cached RDD sits inside a consumer stage, and
following "unpersist cached RDDs" literally can also remove the caches that are
fine. `toDebugString` doesn't help much: it prints the root RDD's storage level
on every line ([RDD.scala
L2072](https://github.com/apache/spark/blob/39ebf2923484b045ef4b4f864db444431503ccd0/core/src/main/scala/org/apache/spark/rdd/RDD.scala#L2072))
and draws pipelined and regular shuffle boundaries the same way, and
`getRDDStorageInfo` skips caches that are not populated yet
([SparkContext.scala
L2399](https://github.com/apache/spark/blob/39ebf2923484b045ef4b4f864db444431503ccd0/core/src/main/scala/org/apache/spark/SparkContext.scala#L2399)).
The repeated-consumption reason at L1598-L1601 alr
eady names the shuffle and the partition. Could this use `find` and include
the offending RDD, e.g. `rdd.toString`, which has the name, class, id and
creation site? With the shared `visited` set, only the first match is reliable,
but that is enough for the message.
##########
core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala:
##########
@@ -1516,6 +1515,29 @@ private[spark] class DAGScheduler(
checkPipelinedStageCreation(finalRDD, partitions, jobId)
}
+ /** Reject persistence that could bypass a native pipelined shuffle reader.
*/
+ private def checkPipelinedCacheSupported(finalRDD: RDD[_]): Unit = {
+ val persisted = new HashSet[RDD[_]]
+ traverseRDDGraph(finalRDD) { (rdd, enqueue) =>
+ if (rdd.getStorageLevel != StorageLevel.NONE ||
Review Comment:
**A cold `localCheckpoint()` on a pipelined consumer is now rejected,
although the base allowed it and, as far as I can tell, handled it safely.**
`localCheckpoint()` persists the RDD ([RDD.scala
L1803](https://github.com/apache/spark/blob/39ebf2923484b045ef4b4f864db444431503ccd0/core/src/main/scala/org/apache/spark/rdd/RDD.scala#L1803)),
so this check rejects a consumer even when none of its blocks exist yet. The
base comment said "Cache / .persist() / local checkpoint are whole-partition
and ephemeral and are not rejected" ([L1488-L1489 at the
base](https://github.com/apache/spark/blob/39ebf2923484b045ef4b4f864db444431503ccd0/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala#L1488-L1489)).
On a cold run, each reader opens once, because pipelined task sets get a
single attempt ([TaskSetManager.scala
L80](https://github.com/apache/spark/blob/39ebf2923484b045ef4b4f864db444431503ccd0/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala#L80))
and repeat
ed reads are rejected. Then `doCheckpoint` runs a job only for the missing
partitions and replaces the lineage with a `LocalCheckpointRDD`
([LocalRDDCheckpointData.scala
L42-L66](https://github.com/apache/spark/blob/39ebf2923484b045ef4b4f864db444431503ccd0/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala#L42-L66)),
so later jobs have no pipelined edge. For example, `c.localCheckpoint();
c.count(); c.map(f).collect()` used to work and now fails at `count()`.
Since persist and reliable checkpoint are rejected too, an RDD-level
pipelined consumer can no longer truncate its lineage inside Spark: an
iterative job re-runs every earlier producer, and its gang demand grows each
iteration. The real hazard needs blocks that already exist at submission, for
example left by a failed run. A narrower rule, such as rejecting only when the
storage level is `NONE` or blocks already exist for the partitions being read,
is possible, but it needs a fresh block-location query and still leaves races,
so I understand keeping the intent rule. If that's the choice, could the
description also say that this removes the in-Spark way to truncate such
lineage?
##########
core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala:
##########
@@ -9470,6 +9471,216 @@ class DAGSchedulerSuite extends SparkFunSuite with
TempLocalSparkContext with Ti
assertDataStructuresEmpty()
}
+ for {
+ (cacheState, cachedPartitions) <- Seq(
+ "cold" -> Seq.empty[Int], "partial" -> Seq(0), "full" -> Seq(0, 1),
+ "evicted" -> Seq(0, 1))
+ slots <- Seq(3, 4)
+ } {
+ test(s"pipelined cache preflight: reject a persisted consumer before
admission " +
+ s"(cache=$cacheState, slots=$slots)") {
+ val producer = new MyRDD(sc, 2, Nil)
+ val dependency = new PipelinedShuffleDependency(producer, new
HashPartitioner(2))
+ val consumer = new MyRDD(sc, 2, List(dependency)).cache()
+ cachedPartitions.foreach { partition =>
+ cacheLocations(consumer.id -> partition) =
Seq(makeBlockManagerId("hostA"))
+ }
+ if (cacheState == "evicted") {
Review Comment:
**The "evicted" case doesn't test a stale cache snapshot, and `slots = 4`
adds nothing over `slots = 3`.** With this PR, all eight cases throw from
`checkPipelinedCacheSupported` (`DAGScheduler` L2693) before admission and
before anything reads `cacheLocs`, so the primed memo is never consulted. On
the base, `handleJobSubmitted` called `clearCacheLocs()` before
`getMissingParentStages`, so the stale snapshot was dropped there too, and
"evicted" behaved like "cold". So the case can't catch an implementation that
"cold" and "full" miss, and its comment ("Keep a stale driver snapshot after
the blocks disappear ...") suggests coverage that isn't there. `slots = 3` is
useful because it pins that the cache check runs before admission. `slots = 4`
only catches a missing check, which "reject persisted consumers with slot
checking disabled" (L9501) also catches. Could we drop "evicted" and `slots =
4`, or reword the comment? That would also save five SparkContext-backed tests.
##########
core/src/test/scala/org/apache/spark/shuffle/local/pipelined/PipelinedChannelShuffleSuite.scala:
##########
@@ -658,6 +659,52 @@ class PipelinedChannelShuffleSuite extends SparkFunSuite {
}
}
+ for (cachedPartitions <- Seq(0, 1, 2)) {
+ test(s"consumer cache is rejected before channel tasks start
(cached=$cachedPartitions)") {
+ withPipelinedSparkContext(cores = 4) { sc =>
+ val input = sc.parallelize(0 until 1000, 2).map(i => (i, i))
+ val consumer = new PipelinedShuffledRDD[Int, Int, Int](input, new
HashPartitioner(2))
+ .cache()
+ val blocks = (0 until cachedPartitions).map(RDDBlockId(consumer.id, _))
+ val blockManager = sc.env.blockManager
+ // Populate blocks without running the shuffle so a partial cache
cannot leave a
+ // producer waiting for a reader. Admission must reject before
starting either stage.
+ blocks.foreach { block =>
+ assert(blockManager.putIterator(
+ block, Iterator((block.splitIndex, block.splitIndex)),
+ StorageLevel.MEMORY_ONLY, tellMaster = true))
+ }
+
assert(blockManager.master.getLocations(blocks.toArray[BlockId]).forall(_.nonEmpty))
+
+ val error = intercept[SparkException](consumer.collect())
+ assert(error.getCondition === "PIPELINED_SHUFFLE_UNSUPPORTED")
+ assert(error.getMessage.contains("unpersist cached RDDs"))
+ assert(sc.statusTracker.getActiveStageIds().isEmpty)
Review Comment:
**These assertions would still pass if the stages had started and the job
had failed later.** `getActiveStageIds()` reads the status store, which the
listener bus updates asynchronously ([SparkStatusTracker.scala
L71-L72](https://github.com/apache/spark/blob/39ebf2923484b045ef4b4f864db444431503ccd0/core/src/main/scala/org/apache/spark/SparkStatusTracker.scala#L71-L72)),
and it is called without `sc.listenerBus.waitUntilEmpty()`. A stage that
started and then failed is not `ACTIVE` either. `numQueuesForTesting` is 0
after any failed run, since `cleanupStateForJobAndIndependentStages` calls
`endRun` (`DAGScheduler` L1817-L1818) before `jobFailed` (`DAGScheduler`
L5311-L5312). And `abortStage` surfaces a user-facing `SparkThrowable`
unwrapped (`DAGScheduler` L5236-L5241), so the condition and message checks
match a late rejection too. A regression that moved the check into
`submitMissingTasks` or into the tasks, with the same reason, would pass this
test. The suite already has `JobEv
ents` (L89-L98) and uses it with `waitUntilEmpty()`. Registering it before
`collect()` and asserting that `taskStarts` and `stageEnds` are empty would
prove what the test name says.
Also, the comment at L670-L671 says "Admission must reject before starting
either stage", but `checkPipelinedCacheSupported` rejects before admission
runs, and a demand of 2 + 2 = 4 tasks fits `local[4]` anyway. For `cached=0`,
the `getLocations(...).forall(_.nonEmpty)` precondition at L677 is vacuous.
##########
core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala:
##########
@@ -9730,7 +9941,17 @@ class DAGSchedulerSuite extends SparkFunSuite with
TempLocalSparkContext with Ti
// Both zip inputs evaluate the same consumer partition. Persistence
cannot guarantee
// that its block survives, and a pending local checkpoint has not
replaced its lineage.
val result = consumer.zipPartitions(consumer)((left, right) => left ++
right)
- assertRepeatedPipelinedConsumptionRejected(result, dependency.shuffleId)
+ if (storage == "uncached") {
+ assertRepeatedPipelinedConsumptionRejected(result,
dependency.shuffleId)
+ } else {
+ assertPipelinedRejectedBeforeStageRegistration(
Review Comment:
**The persisted and local-checkpoint variants of this test no longer test
repeated consumption.** `checkPipelinedCacheSupported` now rejects them before
`checkPipelinedShuffleRepeatedReads` runs, so they only re-check the cache
rejection that "reject a persisted consumer before admission (cache=cold, ...)"
and "reject a pending local checkpoint in a consumer" already cover. Meanwhile
the test name and the comment above ("Both zip inputs evaluate the same
consumer partition ...") still describe the repeated-read check. Could we drop
these two variants, or rename them to say that the cache preflight runs before
the repeated-read check?
Relatedly, the Scaladoc of `checkPipelinedShuffleRepeatedReads`
(`DAGScheduler` L1547, "A cache declaration is not a boundary: evicted or
unstored partitions still compute lineage.") now describes input that function
can no longer receive, because every persisted RDD in `relevantRDDs` is
rejected earlier. It could point to `checkPipelinedCacheSupported` instead, as
the PR already did for the comment at L1488 of that file.
##########
core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala:
##########
@@ -10113,6 +10334,20 @@ class DAGSchedulerSuite extends SparkFunSuite with
TempLocalSparkContext with Ti
}
}
+ private def assertPersistedPipelinedConsumerRejected(
Review Comment:
**This helper is a weaker copy of
`assertPipelinedRejectedBeforeStageRegistration` (L10288), and its
storage-level check is vacuous for two callers.** The existing helper pins the
exact reason with `checkError`, including sqlState `0A000`, submits twice to
show that the rejection leaves no state behind, and checks `results`,
`submittedStageInfos` and `sc.isStopped`. This one regex-matches part of the
reason and submits once without those checks, yet 12 tests use it, while the
self-zip variants at L9947 already use the stronger helper for the same
rejection. Also, in "reject a persisted narrow descendant inside a consumer
stage" (L9517) and "reject a persisted stage that consumes and produces
shuffles" (L9529), the persisted RDD is `cached` or `middle`, not `finalRdd`,
so `finalRdd.getStorageLevel === storageLevel` compares `NONE` with `NONE`.
How about one reason constant in the suite, with this helper as a thin
wrapper like `assertRepeatedPipelinedConsumptionRejected` (L10276): delegate to
the existing helper, then check the storage level of the RDD that was actually
persisted? Today the reason is copied verbatim at L9949-L9952 and pinned by
different substrings at L9551, L9553, L9643, L10341-L10344 and in
`PipelinedChannelShuffleSuite` L681. That last one also checks the condition by
hand instead of using `checkError`, unlike L249-L255 of the same suite.
##########
core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala:
##########
@@ -1516,6 +1515,29 @@ private[spark] class DAGScheduler(
checkPipelinedStageCreation(finalRDD, partitions, jobId)
}
+ /** Reject persistence that could bypass a native pipelined shuffle reader.
*/
+ private def checkPipelinedCacheSupported(finalRDD: RDD[_]): Unit = {
+ val persisted = new HashSet[RDD[_]]
+ traverseRDDGraph(finalRDD) { (rdd, enqueue) =>
+ if (rdd.getStorageLevel != StorageLevel.NONE ||
+
rdd.checkpointData.exists(_.isInstanceOf[LocalRDDCheckpointData[_]])) {
+ persisted += rdd
+ }
+ rdd.dependencies.foreach(dep => enqueue(dep.rdd))
Review Comment:
**This walk also descends into the materialized prefix, where a pipelined
read can't exist.** `handleJobSubmitted` has already rejected a pipelined
shuffle below a regular boundary (`isUnsupportedMix`, L2677) when this runs, so
no RDD that is reachable only through a regular `ShuffleDependency` can have a
within-stage pipelined read. The walk still visits the whole prefix and then
searches from every persisted RDD in it. In an RDD job with a materialized
prefix, that prefix is often the long, cached part of an iterative program, and
this runs on the event loop for every pipelined submission. The updated tests
show the extra walk: 2 -> 3 and 3/4 -> 4/5 visits per edge. Stopping at regular
boundaries keeps the result and drops that cost:
```scala
rdd.dependencies.foreach {
case pd: PipelinedShuffleDependency[_, _, _] => enqueue(pd.rdd)
case _: ShuffleDependency[_, _, _] => // materialized prefix, excluded by
isUnsupportedMix
case dep => enqueue(dep.rdd)
}
```
A larger option is to collect the persisted RDDs in
`checkPipelinedJobSupported`'s walk (L1440-L1469), moved ahead of admission,
and test membership in `relevantRDDs` (L1555-L1563), the narrow-descendant
closure of the pipelined readers. The demand walk can't host it, because it is
skipped when the slot check is off.
##########
core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala:
##########
@@ -9470,6 +9471,216 @@ class DAGSchedulerSuite extends SparkFunSuite with
TempLocalSparkContext with Ti
assertDataStructuresEmpty()
}
+ for {
+ (cacheState, cachedPartitions) <- Seq(
+ "cold" -> Seq.empty[Int], "partial" -> Seq(0), "full" -> Seq(0, 1),
+ "evicted" -> Seq(0, 1))
+ slots <- Seq(3, 4)
+ } {
+ test(s"pipelined cache preflight: reject a persisted consumer before
admission " +
+ s"(cache=$cacheState, slots=$slots)") {
+ val producer = new MyRDD(sc, 2, Nil)
+ val dependency = new PipelinedShuffleDependency(producer, new
HashPartitioner(2))
+ val consumer = new MyRDD(sc, 2, List(dependency)).cache()
+ cachedPartitions.foreach { partition =>
+ cacheLocations(consumer.id -> partition) =
Seq(makeBlockManagerId("hostA"))
+ }
+ if (cacheState == "evicted") {
+ // Keep a stale driver snapshot after the blocks disappear. Cache
presence cannot
+ // authorize execution of a transient shuffle whose readers can be
bypassed.
+ assert(scheduler.getCacheLocs(consumer).forall(_.nonEmpty))
+ cacheLocations.clear()
+ assert(scheduler.getCacheLocs(consumer).forall(_.nonEmpty))
+ }
+ scheduler.asInstanceOf[MyDAGScheduler].maxConcurrentTasksForTest = slots
+
+ assertPersistedPipelinedConsumerRejected(consumer,
Seq(dependency.shuffleId))
+ }
+ }
+
+ test("pipelined cache preflight: reject persisted consumers with slot
checking disabled") {
+ conf.set(config.PIPELINED_GROUP_SLOT_CHECK_ENABLED, false)
+ val dependency = new PipelinedShuffleDependency(new MyRDD(sc, 2, Nil), new
HashPartitioner(2))
+ val consumer = new MyRDD(sc, 2,
List(dependency)).persist(StorageLevel.MEMORY_AND_DISK)
+ cacheLocations(consumer.id -> 0) = Seq(makeBlockManagerId("hostA"))
+
+ assertPersistedPipelinedConsumerRejected(consumer,
Seq(dependency.shuffleId))
+ }
+
+ test("pipelined cache preflight: reject a persisted narrow descendant inside
a consumer stage") {
+ val dependency = new PipelinedShuffleDependency(new MyRDD(sc, 2, Nil), new
HashPartitioner(2))
+ val reads = new MyRDD(sc, 2, List(dependency))
+ val cached = new MyRDD(sc, 2, List(new OneToOneDependency(reads))).cache()
+ val result = new MyRDD(sc, 2, List(new OneToOneDependency(cached)))
+ cacheLocations(cached.id -> 0) = Seq(makeBlockManagerId("hostA"))
+
+ assertPersistedPipelinedConsumerRejected(result, Seq(dependency.shuffleId))
+ }
+
+ test("pipelined cache preflight: reject a persisted stage that consumes and
produces shuffles") {
+ val first = new PipelinedShuffleDependency(new MyRDD(sc, 2, Nil), new
HashPartitioner(2))
+ val middle = new MyRDD(sc, 2, List(first)).cache()
+ val second = new PipelinedShuffleDependency(middle, new HashPartitioner(2))
+ val result = new MyRDD(sc, 2, List(second))
+ for (partition <- 0 until 2) {
+ cacheLocations(middle.id -> partition) = Seq(makeBlockManagerId("hostA"))
+ }
+
+ assertPersistedPipelinedConsumerRejected(result, Seq(first.shuffleId,
second.shuffleId))
+ }
+
+ test("pipelined cache preflight: reject a pending local checkpoint in a
consumer") {
+ val dependency = new PipelinedShuffleDependency(new MyRDD(sc, 2, Nil), new
HashPartitioner(2))
+ val consumer = new MyRDD(sc, 2, List(dependency)).localCheckpoint()
+ assert(!consumer.isCheckpointed)
+ assert(consumer.getStorageLevel != StorageLevel.NONE)
+
+ assertPersistedPipelinedConsumerRejected(consumer,
Seq(dependency.shuffleId))
+ }
+
+ test("pipelined cache preflight: unpersisting a pending local checkpoint
still requires " +
+ "rebuilding the consumer") {
+ val dependency = new PipelinedShuffleDependency(new MyRDD(sc, 2, Nil), new
HashPartitioner(2))
+ val consumer = new MyRDD(sc, 2, List(dependency)).localCheckpoint()
+ consumer.unpersist(blocking = true)
+ assert(consumer.getStorageLevel === StorageLevel.NONE)
+
assert(consumer.checkpointData.exists(_.isInstanceOf[LocalRDDCheckpointData[_]]))
+ assert(!consumer.isCheckpointed)
+
+ val exception = submitAndCaptureFailure(consumer, Array(0, 1))
+ assertPipelinedUnsupported(exception, "cache/persist/localCheckpoint in a
pipelined consumer")
+ assert(exception.getMessage.contains(
+ "rebuild locally checkpointed consumers without localCheckpoint"))
+ assert(consumer.getStorageLevel === StorageLevel.NONE)
+ assertNoPipelinedStageRegistration(Seq(dependency.shuffleId))
+ assert(taskSets.isEmpty)
+ assertDataStructuresEmpty()
+ }
+
+ for (cachedProducerRoot <- Seq(false, true)) {
+ test(s"pipelined cache preflight: allow cached producer input " +
+ s"(producerRoot=$cachedProducerRoot)") {
+ val input = new MyRDD(sc, 2, Nil)
+ val producer = new MyRDD(sc, 2, List(new OneToOneDependency(input)))
+ val cached = if (cachedProducerRoot) producer else input
+ cached.cache()
+ for (partition <- 0 until 2) {
+ cacheLocations(cached.id -> partition) =
Seq(makeBlockManagerId("hostA"))
+ }
+ val dependency = new PipelinedShuffleDependency(producer, new
HashPartitioner(2))
+ val result = new MyRDD(sc, 2, List(dependency))
+ scheduler.asInstanceOf[MyDAGScheduler].maxConcurrentTasksForTest = 4
+
+ submit(result, Array(0, 1))
+ assert(failure === null)
+ assert(taskSets.size === 2)
+ assert(taskSets.forall(_.isPipelined))
+ completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2)
+ complete(taskSets.last, Seq((Success, 42), (Success, 43)))
+ assert(results === Map(0 -> 42, 1 -> 43))
+ assertDataStructuresEmpty()
+ }
+ }
+
+ test("pipelined cache preflight: allow unrelated cached branches with shared
ancestry") {
+ val edgeReads = Seq.fill(8)(new AtomicInteger)
+ var shared: RDD[(Int, Int)] = new MyRDD(sc, 2, Nil)
+ edgeReads.foreach { counter =>
+ val dependency = new OneToOneDependency(shared) {
+ override def rdd: RDD[(Int, Int)] = {
+ counter.incrementAndGet()
+ super.rdd
+ }
+ }
+ shared = new MyRDD(sc, 2, List(dependency))
+ }
+ val cachedBranches = (0 until 32).map { _ =>
+ val cached = new MyRDD(sc, 2, List(new
OneToOneDependency(shared))).cache()
+ cacheLocations(cached.id -> 0) = Seq(makeBlockManagerId("hostA"))
+ cached
+ }
+ val pipelined = new PipelinedShuffleDependency(new MyRDD(sc, 2, Nil), new
HashPartitioner(2))
+ val reads = new MyRDD(sc, 2, List(pipelined))
+ val result = new MyRDD(sc, 2,
+ (reads +: cachedBranches).map(new OneToOneDependency(_)).toList)
+ var preflightEdgeReads = Seq.empty[Int]
+ doAnswer { invocation =>
+ if (preflightEdgeReads.isEmpty) {
+ preflightEdgeReads = edgeReads.map(_.get())
+ }
+ invocation.callRealMethod()
+ }.when(scheduler).getShuffleDependenciesAndResourceProfiles(result)
+ edgeReads.foreach(_.set(0))
+
+ submit(result, Array(0, 1))
+ assert(failure === null)
+ // Every cached branch shares a safe ancestor. Searching each branch
independently would
+ // visit each edge at least 32 times before stage creation.
+ assert(preflightEdgeReads.nonEmpty)
+ assert(preflightEdgeReads.forall(_ <= 6), s"repeated shared traversal:
$preflightEdgeReads")
Review Comment:
**`<= 6` is looser than the actual count, which is exactly 5.** Before the
first `getShuffleDependenciesAndResourceProfiles(result)` call, each counted
edge is read by the classification walk, the cache collection walk, the shared
cache search, the admission demand walk and `checkPipelinedJobSupported`'s
walk. So a regression that adds one more persisted-only search with its own
`visited` set (5 -> 6) still passes, and the nearby test that pins exact counts
(L9797-L9798) has no persisted RDD. The setup also repeats about 20 lines of
"pipelined checkpoint preflight: shared ancestry is visited once" (L9746).
Could this be a `"cache"` value of that test's `roots` dimension, with
`expected = 5`, so the count is pinned exactly?
##########
core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala:
##########
@@ -1516,6 +1515,29 @@ private[spark] class DAGScheduler(
checkPipelinedStageCreation(finalRDD, partitions, jobId)
}
+ /** Reject persistence that could bypass a native pipelined shuffle reader.
*/
+ private def checkPipelinedCacheSupported(finalRDD: RDD[_]): Unit = {
+ val persisted = new HashSet[RDD[_]]
+ traverseRDDGraph(finalRDD) { (rdd, enqueue) =>
+ if (rdd.getStorageLevel != StorageLevel.NONE ||
+
rdd.checkpointData.exists(_.isInstanceOf[LocalRDDCheckpointData[_]])) {
Review Comment:
**nit: `RDD.isLocallyCheckpointed` already implements this predicate.** It
is the check that `persist` and `unpersist` use for "marked for local
checkpointing" ([RDD.scala
L1849-L1854](https://github.com/apache/spark/blob/39ebf2923484b045ef4b4f864db444431503ccd0/core/src/main/scala/org/apache/spark/rdd/RDD.scala#L1849-L1854),
used at
[L187](https://github.com/apache/spark/blob/39ebf2923484b045ef4b4f864db444431503ccd0/core/src/main/scala/org/apache/spark/rdd/RDD.scala#L187)
and
[L214](https://github.com/apache/spark/blob/39ebf2923484b045ef4b4f864db444431503ccd0/core/src/main/scala/org/apache/spark/rdd/RDD.scala#L214)),
and like this expression it is true for both pending and materialized local
checkpoints. Widening it from `private[rdd]` to `private[spark]`, as
`isCheckpointedAndMaterialized` and `isReliablyCheckpointed` already are
([L1842](https://github.com/apache/spark/blob/39ebf2923484b045ef4b4f864db444431503ccd0/core/src/main/scala/org/apache/spark/rdd/RDD.scala#L1842),
[L185
9](https://github.com/apache/spark/blob/39ebf2923484b045ef4b4f864db444431503ccd0/core/src/main/scala/org/apache/spark/rdd/RDD.scala#L1859)),
would keep a single definition, and both new `LocalRDDCheckpointData` imports
could go away (the suite uses it only at L9547). The reliable check at L1441
has to stay spelled out, since `isReliablyCheckpointed` is true only after
materialization.
##########
core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala:
##########
@@ -2665,19 +2687,17 @@ private[spark] class DAGScheduler(
return
}
- // Gang admission for a pipelined job: the runnable stage graph is one
pipelined group (a
- // materialized prefix's stages never run), 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)) {
- // Clear the slot-check retry count an earlier deferral may have left.
- barrierJobIdToNumTasksCheckFailures.remove(jobId)
- return
- }
var finalStage: ResultStage = null
try {
+ if (hasPipelined) {
+ checkPipelinedCacheSupported(finalRDD)
+ // Validate cache semantics before counting tasks: a cached consumer
could suppress
Review Comment:
**nit: this comment explains the call above it, and the two `if
(hasPipelined)` blocks could be one.** Since it sits between
`checkPipelinedCacheSupported(finalRDD)` and the admission call, it reads as a
comment on admission. The moved
`barrierJobIdToNumTasksCheckFailures.remove(jobId)` also lost its "Clear the
slot-check retry count an earlier deferral may have left." comment, so it is
the only one of the six without it. Merging the blocks (cache check, admission,
`checkPipelinedJobSupported`) and moving this comment above the cache check
would read better.
Moving admission into the `try` also changes behavior a little. A failure
inside admission that isn't an idiom rejection, such as an invalid
`spark.scheduler.pipelinedGroup.slotCheck.enabled` value, now fails only this
job through the generic `case e: Exception` arm, instead of reaching `onError`
and stopping the SparkContext. That looks like an improvement. Could the
description mention it?
##########
core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala:
##########
@@ -9470,6 +9471,216 @@ class DAGSchedulerSuite extends SparkFunSuite with
TempLocalSparkContext with Ti
assertDataStructuresEmpty()
}
+ for {
+ (cacheState, cachedPartitions) <- Seq(
+ "cold" -> Seq.empty[Int], "partial" -> Seq(0), "full" -> Seq(0, 1),
+ "evicted" -> Seq(0, 1))
+ slots <- Seq(3, 4)
+ } {
+ test(s"pipelined cache preflight: reject a persisted consumer before
admission " +
+ s"(cache=$cacheState, slots=$slots)") {
+ val producer = new MyRDD(sc, 2, Nil)
+ val dependency = new PipelinedShuffleDependency(producer, new
HashPartitioner(2))
+ val consumer = new MyRDD(sc, 2, List(dependency)).cache()
+ cachedPartitions.foreach { partition =>
+ cacheLocations(consumer.id -> partition) =
Seq(makeBlockManagerId("hostA"))
+ }
+ if (cacheState == "evicted") {
+ // Keep a stale driver snapshot after the blocks disappear. Cache
presence cannot
+ // authorize execution of a transient shuffle whose readers can be
bypassed.
+ assert(scheduler.getCacheLocs(consumer).forall(_.nonEmpty))
+ cacheLocations.clear()
+ assert(scheduler.getCacheLocs(consumer).forall(_.nonEmpty))
+ }
+ scheduler.asInstanceOf[MyDAGScheduler].maxConcurrentTasksForTest = slots
+
+ assertPersistedPipelinedConsumerRejected(consumer,
Seq(dependency.shuffleId))
+ }
+ }
+
+ test("pipelined cache preflight: reject persisted consumers with slot
checking disabled") {
+ conf.set(config.PIPELINED_GROUP_SLOT_CHECK_ENABLED, false)
+ val dependency = new PipelinedShuffleDependency(new MyRDD(sc, 2, Nil), new
HashPartitioner(2))
+ val consumer = new MyRDD(sc, 2,
List(dependency)).persist(StorageLevel.MEMORY_AND_DISK)
+ cacheLocations(consumer.id -> 0) = Seq(makeBlockManagerId("hostA"))
+
+ assertPersistedPipelinedConsumerRejected(consumer,
Seq(dependency.shuffleId))
+ }
+
+ test("pipelined cache preflight: reject a persisted narrow descendant inside
a consumer stage") {
+ val dependency = new PipelinedShuffleDependency(new MyRDD(sc, 2, Nil), new
HashPartitioner(2))
+ val reads = new MyRDD(sc, 2, List(dependency))
+ val cached = new MyRDD(sc, 2, List(new OneToOneDependency(reads))).cache()
+ val result = new MyRDD(sc, 2, List(new OneToOneDependency(cached)))
+ cacheLocations(cached.id -> 0) = Seq(makeBlockManagerId("hostA"))
+
+ assertPersistedPipelinedConsumerRejected(result, Seq(dependency.shuffleId))
+ }
+
+ test("pipelined cache preflight: reject a persisted stage that consumes and
produces shuffles") {
+ val first = new PipelinedShuffleDependency(new MyRDD(sc, 2, Nil), new
HashPartitioner(2))
+ val middle = new MyRDD(sc, 2, List(first)).cache()
+ val second = new PipelinedShuffleDependency(middle, new HashPartitioner(2))
+ val result = new MyRDD(sc, 2, List(second))
+ for (partition <- 0 until 2) {
+ cacheLocations(middle.id -> partition) = Seq(makeBlockManagerId("hostA"))
+ }
+
+ assertPersistedPipelinedConsumerRejected(result, Seq(first.shuffleId,
second.shuffleId))
+ }
+
+ test("pipelined cache preflight: reject a pending local checkpoint in a
consumer") {
+ val dependency = new PipelinedShuffleDependency(new MyRDD(sc, 2, Nil), new
HashPartitioner(2))
+ val consumer = new MyRDD(sc, 2, List(dependency)).localCheckpoint()
+ assert(!consumer.isCheckpointed)
+ assert(consumer.getStorageLevel != StorageLevel.NONE)
+
+ assertPersistedPipelinedConsumerRejected(consumer,
Seq(dependency.shuffleId))
+ }
+
+ test("pipelined cache preflight: unpersisting a pending local checkpoint
still requires " +
+ "rebuilding the consumer") {
+ val dependency = new PipelinedShuffleDependency(new MyRDD(sc, 2, Nil), new
HashPartitioner(2))
+ val consumer = new MyRDD(sc, 2, List(dependency)).localCheckpoint()
+ consumer.unpersist(blocking = true)
+ assert(consumer.getStorageLevel === StorageLevel.NONE)
+
assert(consumer.checkpointData.exists(_.isInstanceOf[LocalRDDCheckpointData[_]]))
+ assert(!consumer.isCheckpointed)
+
+ val exception = submitAndCaptureFailure(consumer, Array(0, 1))
+ assertPipelinedUnsupported(exception, "cache/persist/localCheckpoint in a
pipelined consumer")
+ assert(exception.getMessage.contains(
+ "rebuild locally checkpointed consumers without localCheckpoint"))
+ assert(consumer.getStorageLevel === StorageLevel.NONE)
+ assertNoPipelinedStageRegistration(Seq(dependency.shuffleId))
+ assert(taskSets.isEmpty)
+ assertDataStructuresEmpty()
+ }
+
+ for (cachedProducerRoot <- Seq(false, true)) {
+ test(s"pipelined cache preflight: allow cached producer input " +
+ s"(producerRoot=$cachedProducerRoot)") {
+ val input = new MyRDD(sc, 2, Nil)
+ val producer = new MyRDD(sc, 2, List(new OneToOneDependency(input)))
+ val cached = if (cachedProducerRoot) producer else input
+ cached.cache()
+ for (partition <- 0 until 2) {
+ cacheLocations(cached.id -> partition) =
Seq(makeBlockManagerId("hostA"))
+ }
+ val dependency = new PipelinedShuffleDependency(producer, new
HashPartitioner(2))
+ val result = new MyRDD(sc, 2, List(dependency))
+ scheduler.asInstanceOf[MyDAGScheduler].maxConcurrentTasksForTest = 4
+
+ submit(result, Array(0, 1))
+ assert(failure === null)
+ assert(taskSets.size === 2)
+ assert(taskSets.forall(_.isPipelined))
+ completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2)
+ complete(taskSets.last, Seq((Success, 42), (Success, 43)))
+ assert(results === Map(0 -> 42, 1 -> 43))
+ assertDataStructuresEmpty()
+ }
+ }
+
+ test("pipelined cache preflight: allow unrelated cached branches with shared
ancestry") {
+ val edgeReads = Seq.fill(8)(new AtomicInteger)
+ var shared: RDD[(Int, Int)] = new MyRDD(sc, 2, Nil)
+ edgeReads.foreach { counter =>
+ val dependency = new OneToOneDependency(shared) {
+ override def rdd: RDD[(Int, Int)] = {
+ counter.incrementAndGet()
+ super.rdd
+ }
+ }
+ shared = new MyRDD(sc, 2, List(dependency))
+ }
+ val cachedBranches = (0 until 32).map { _ =>
+ val cached = new MyRDD(sc, 2, List(new
OneToOneDependency(shared))).cache()
+ cacheLocations(cached.id -> 0) = Seq(makeBlockManagerId("hostA"))
+ cached
+ }
+ val pipelined = new PipelinedShuffleDependency(new MyRDD(sc, 2, Nil), new
HashPartitioner(2))
+ val reads = new MyRDD(sc, 2, List(pipelined))
+ val result = new MyRDD(sc, 2,
+ (reads +: cachedBranches).map(new OneToOneDependency(_)).toList)
+ var preflightEdgeReads = Seq.empty[Int]
+ doAnswer { invocation =>
+ if (preflightEdgeReads.isEmpty) {
+ preflightEdgeReads = edgeReads.map(_.get())
+ }
+ invocation.callRealMethod()
+ }.when(scheduler).getShuffleDependenciesAndResourceProfiles(result)
+ edgeReads.foreach(_.set(0))
+
+ submit(result, Array(0, 1))
+ assert(failure === null)
+ // Every cached branch shares a safe ancestor. Searching each branch
independently would
+ // visit each edge at least 32 times before stage creation.
+ assert(preflightEdgeReads.nonEmpty)
+ assert(preflightEdgeReads.forall(_ <= 6), s"repeated shared traversal:
$preflightEdgeReads")
+ assert(taskSets.size === 2)
+ completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2)
+ complete(taskSets.last, Seq((Success, 42), (Success, 43)))
+ assert(results === Map(0 -> 42, 1 -> 43))
+ assertDataStructuresEmpty()
+ }
+
+ test("pipelined cache preflight: persistence changes are checked on every
submission") {
+ val dependency = new PipelinedShuffleDependency(new MyRDD(sc, 2, Nil), new
HashPartitioner(2))
+ val consumer = new MyRDD(sc, 2, List(dependency))
+ submit(consumer, Array(0, 1))
+ completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2)
+ complete(taskSets.last, Seq((Success, 42), (Success, 43)))
+ assertDataStructuresEmpty()
+
+ consumer.persist(StorageLevel.MEMORY_AND_DISK)
+ for (partition <- 0 until 2) {
+ cacheLocations(consumer.id -> partition) =
Seq(makeBlockManagerId("hostA"))
+ }
+ clearInvocations(scheduler, mapOutputTracker)
+ assertPipelinedUnsupported(
+ submitAndCaptureFailure(consumer, Array(0, 1)),
+ "cache/persist/localCheckpoint in a pipelined consumer")
+ assertNoPipelinedStageRegistration(Nil)
+ assert(taskSets.size === 2, "the rejected submission must launch no
additional tasks")
+ assertDataStructuresEmpty()
+
+ consumer.unpersist(blocking = true)
+ assert(consumer.getStorageLevel === StorageLevel.NONE)
+ results.clear()
+ submit(consumer, Array(0, 1))
+ assert(taskSets.size === 4)
+ assert(failure === null)
+ completeShuffleMapStageSuccessfully(taskSets(2).stageId, 0, 2)
+ complete(taskSets.last, Seq((Success, 44), (Success, 45)))
+ assert(results === Map(0 -> 44, 1 -> 45))
+ assertDataStructuresEmpty()
+ }
+
+ for (cachedPartitions <- Seq(Seq.empty[Int], Seq(0), Seq(0, 1))) {
Review Comment:
**nit: these three tests can't fail because of this PR.** They use a regular
`ShuffleDependency`, so `hasPipelined` is false and the new check never runs.
Even without that gate, `rddChainReadsPipelinedShuffle` stops at the regular
boundary, and admission would see a demand of 2 against
`maxConcurrentTasksForTest = 1000`. Cache pruning for regular jobs is already
covered by "getMissingParentStages should consider all ancestor RDDs' cache
statuses" (L899) and "SPARK-53564: Resubmit missing parent when failed to get
cache locations" (L6748), and the `hasPipelined` gate is pinned by "ordinary
job preflight traverses once" (L7226). Could we drop them, or keep just one?
##########
core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala:
##########
@@ -5554,12 +5574,13 @@ private[spark] object DAGScheduler {
private[spark] case class CancelledJobInfo(jobId: Int, properties: Properties)
/**
- * 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,
+ * Thrown when a job uses a pipelined-shuffle idiom that is not supported
(fan-out, dynamic
Review Comment:
**nit: "dynamic allocation" is listed here, but a result job with dynamic
allocation gets a plain `SparkException` without this condition.**
`classifyJobShuffleShapeAndCheckFeatures` rejects it first (L1370-L1373), and
`DAGSchedulerSuite` L7394 pins that message, so the typed branch at L861-L862
is reached only through the map-stage path. Meanwhile, repeated consumption
(L1598-L1601) throws this type but is still not listed. Could the list qualify
or drop the dynamic-allocation item and add repeated consumption?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
