dongjoon-hyun commented on code in PR #59523:
URL: https://github.com/apache/spark/pull/59523#discussion_r4426582640
##########
core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala:
##########
@@ -1520,13 +1519,39 @@ 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.isLocallyCheckpointed) {
+ persisted += rdd
+ }
+ rdd.dependencies.foreach {
+ case pd: PipelinedShuffleDependency[_, _, _] => enqueue(pd.rdd)
+ case _: ShuffleDependency[_, _, _] =>
+ // Shape validation excludes pipelined readers below a regular
boundary.
+ case 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[_]]
+ persisted.find(rdd => rddChainReadsPipelinedShuffle(rdd, visited)).foreach
{ rdd =>
Review Comment:
**Follow-up to my item 2: when several persisted RDDs read the pipelined
shuffle, which one the error names can change from run to run.** `persisted` is
a `HashSet[RDD[_]]` (L1524), and `RDD` doesn't override `hashCode`, so `find`
checks the roots in an order based on identity hash codes, which differ between
runs. For example, with `val c = consumer.cache()`, where `consumer` reads the
pipelined shuffle, and `val d = c.map(identity).cache()`, `d.count()` can name
either `c` or `d`. The same job can then log different errors across runs, and
a test with two offenders can't pin the name. `traverseRDDGraph` already calls
the visitor once per RDD, so a buffer would hold no duplicates and keep the
traversal order. Could `persisted` be a `ListBuffer`, which is already
imported, so the error always names the first offender in that order?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/PipelinedShuffleEligibility.scala:
##########
@@ -49,6 +50,9 @@ private[sql] object PipelinedShuffleEligibility extends
Logging {
def isUnsupportedConsumer(plan: SparkPlan): Boolean = plan match {
case _: CoalesceExec | _: CartesianProductExec | _:
DeserializeToObjectExec |
_: CollectLimitExec | _: CollectTailExec | _:
TakeOrderedAndProjectExec => true
+ // Caching can persist a shuffle reader before zipWithIndex submits its
counting job.
Review Comment:
**The Scaladoc at L49 and the class doc's lead-in at
`EnablePipelinedShuffle.scala:51-53` don't cover this new reason.** L49 still
describes these as "Operators whose consumers cannot safely drain a bounded,
single-reader channel", and `EnablePipelinedShuffle.scala:51-53` describes an
unsupported consumer as one that drains a shuffle in a way the channel
transport cannot serve, or builds its own hidden regular shuffle. A cached
`AttachDistributedSequenceExec` does neither. It is unsupported because
persisting its child can skip channel readers, as this comment and the new
bullet at `EnablePipelinedShuffle.scala:75-78` say. Both places are outside
this diff, so I left this here. Could they mention persistence too, e.g. "...
or that can persist a channel reader"?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/PipelinedShuffleSqlSuite.scala:
##########
@@ -213,7 +218,89 @@ class PipelinedShuffleSqlSuite extends SparkFunSuite
}
}
+ private def distributedSequence(
Review Comment:
**This helper builds the logical `AttachDistributedSequence` directly, so
these tests skip the path pandas-on-Spark takes.** pandas-on-Spark attaches its
default index with `sdf.select(SF.distributed_sequence_id().alias(column_name),
"*")` ([internal.py
L952](https://github.com/apache/spark/blob/8622bf64d253128af42121f0da257c055961a627/python/pyspark/pandas/internal.py#L952)),
where `SF.distributed_sequence_id()` passes `F.lit(True)` as the `cache`
argument ([pyspark/sql/internal.py
L108-L111](https://github.com/apache/spark/blob/8622bf64d253128af42121f0da257c055961a627/python/pyspark/sql/internal.py#L108-L111)).
`ExtractDistributedSequenceID` then reads that flag and puts a `Project` above
the operator ([ExtractDistributedSequenceID.scala
L37-L44](https://github.com/apache/spark/blob/8622bf64d253128af42121f0da257c055961a627/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ExtractDistributedSequenceID.scala#L37-L44)).
Could the helper build the same `select` with
`Column.internalFn("distributed_sequence_id", lit(cache))`, as
`DataFrameSuite` does at
[L2422](https://github.com/apache/spark/blob/8622bf64d253128af42121f0da257c055961a627/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala#L2422)
without the argument? That would also drop the `classic` cast, and since the
description says the Python pandas wrapper was not executed, it would cover the
plan that wrapper builds.
##########
core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala:
##########
@@ -9499,6 +9554,146 @@ class DAGSchedulerSuite extends SparkFunSuite with
TempLocalSparkContext with Ti
assertDataStructuresEmpty()
}
+ for ((cacheState, cachedPartitions) <- Seq(
+ "cold" -> Seq.empty[Int], "partial" -> Seq(0), "full" -> Seq(0, 1))) {
+ test(s"pipelined cache preflight: reject a persisted consumer before
admission " +
+ s"(cache=$cacheState)") {
+ 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"))
+ }
+ scheduler.asInstanceOf[MyDAGScheduler].maxConcurrentTasksForTest = 3
+
+ assertPersistedPipelinedConsumerRejected(consumer, 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"))
Review Comment:
**nit: this cache location is never read.** `checkPipelinedCacheSupported`
rejects the job before admission and stage creation, and only
`getMissingParentStages` and the preferred-location lookup call `getCacheLocs`,
so the test passes without this line, like the "evicted" snapshot in my item 5.
The cold/partial/full loop above already covers cache state. Could we drop it,
so this test only varies the slot check?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnablePipelinedShuffle.scala:
##########
@@ -127,24 +131,26 @@ object EnablePipelinedShuffle extends Rule[SparkPlan] {
* True if any [[ShuffleExchangeExec]] in `plan` is read by an operator the
channel transport
* cannot serve. The unsupported operators (see class doc for why each is
fatal) are
* `CoalesceExec`, `CartesianProductExec`, and the limit operators
`CollectLimitExec` /
- * `CollectTailExec` / `TakeOrderedAndProjectExec`. For each such operator
anywhere in the plan,
- * check whether a shuffle is reachable below it.
- *
- * The reachability walk descends through EVERY child of a non-exchange node
-- not only unary
- * children -- so a shuffle behind a `UnionExec`/join (a `BinaryExecNode`)
beneath the operator
- * is still found. It stops at the FIRST [[ShuffleExchangeExec]] on each
path: a shuffle deeper
- * than that first one is not read by this operator (the intervening
exchange's own reader reads
- * one reduce partition per task), so it is not this operator's concern.
+ * `CollectTailExec` / `TakeOrderedAndProjectExec`, plus cached
`AttachDistributedSequenceExec`.
+ * A shuffle above an unsupported operator with shuffle-free input does not
disqualify the plan.
*/
- private def readsShuffleThroughUnsupportedConsumer(plan: SparkPlan): Boolean
= {
- def reachesShuffle(p: SparkPlan): Boolean = p match {
- case _: ShuffleExchangeExec => true
- case other => other.children.exists(reachesShuffle)
- }
- plan.exists {
- case p if PipelinedShuffleEligibility.isUnsupportedConsumer(p) =>
- p.children.exists(reachesShuffle)
- case _ => false
+ private[exchange] def readsShuffleThroughUnsupportedConsumer(plan:
SparkPlan): Boolean = {
Review Comment:
**This rewrite replaces an 8-line `plan.exists` with shared mutable state,
to save planning work that seems small.** The old version was quadratic only
when unsupported consumers nest with no exchange between them, and even at the
new grid test's depth of 128 that is about 16K node visits during planning. In
exchange, the new version has a `foundUnsupported` flag shared across the
recursion, a manual iterator loop for the early exit, and a grid test with a
custom `CountingChildrenExec` to pin the visit count. Could we go back to the
`plan.exists` version? If the linear walk is needed, a post-order fold that
returns `(hasShuffle, found)` would avoid the shared `var`.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/PipelinedShuffleSqlSuite.scala:
##########
@@ -213,7 +218,89 @@ class PipelinedShuffleSqlSuite extends SparkFunSuite
}
}
+ private def distributedSequence(
+ spark: SparkSession,
+ input: DataFrame,
+ cache: Boolean): DataFrame = {
+ classic.Dataset.ofRows(spark.asInstanceOf[classic.SparkSession],
+ AttachDistributedSequence(AttributeReference("index", LongType)(),
+ input.queryExecution.analyzed, cache))
+ }
+
+ private def assertDistributedSequenceResult(df: DataFrame): Unit = {
+ val rows = df.collect()
+ assert(rows.map(_.getLong(0)).sorted === (0L until 4000L).toArray)
+ assert(rows.map(_.getLong(1)).sorted === (0L until 4000L).toArray)
+ }
+
for (aqe <- Seq(false, true)) {
+ test(s"distributed sequence cache construction retains regular shuffles
with AQE=$aqe") {
+ withPipelinedSession("pipelined-distributed-sequence-cache", aqe,
+ channelBatchSize = 8, channelQueueCapacity = 1) { spark =>
+ val cacheKey = "pandas_on_Spark.compute.default_index_cache"
+ for {
+ storage <- Seq(None, Some("MEMORY_ONLY"), Some("LOCAL_CHECKPOINT"),
Some("NONE"))
+ upperShuffle <- Seq(false, true)
+ } {
+ // An absent setting exercises pandas-on-Spark's default
MEMORY_AND_DISK_SER.
+ storage match {
+ case Some(level) => spark.conf.set(cacheKey, "\"" + level + "\"")
+ case None => spark.conf.unset(cacheKey)
+ }
+ val indexed = distributedSequence(spark,
+ spark.range(0, 4000, 1, 2).repartition(4).toDF(), cache = true)
+ val df = if (upperShuffle) indexed.repartition(3) else indexed
Review Comment:
**With AQE on, the `upperShuffle = true` iterations pass without the new
`isUnsupportedConsumer` case.** `collectCandidates` collects the
`repartition(3)` exchange and does not walk below it, since it only continues
below a `SinglePartition` exchange ([AQEEnablePipelinedShuffle.scala
L118-L121](https://github.com/apache/spark/blob/8622bf64d253128af42121f0da257c055961a627/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEEnablePipelinedShuffle.scala#L118-L121)).
The `repartition(4)` exchange is never a candidate, so `hasRegularPrefix`
keeps the whole plan regular ([AQEEnablePipelinedShuffle.scala
L72-L79](https://github.com/apache/spark/blob/8622bf64d253128af42121f0da257c055961a627/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEEnablePipelinedShuffle.scala#L72-L79))
whether or not the cached operator is an unsupported consumer. So under AQE,
only the `upperShuffle = false` iterations check the new case. Could the upper
shuffle be `repartition(1)
`? `Repartition` plans one partition as `SinglePartition`
([basicLogicalOperators.scala
L2285-L2289](https://github.com/apache/spark/blob/8622bf64d253128af42121f0da257c055961a627/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala#L2285-L2289)),
so the walk would reach the cached operator, where `isStatsSensitive` now
blocks the `repartition(4)` exchange ([AQEEnablePipelinedShuffle.scala
L162](https://github.com/apache/spark/blob/8622bf64d253128af42121f0da257c055961a627/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEEnablePipelinedShuffle.scala#L162)),
and the iteration would fail without the new case.
##########
core/src/test/scala/org/apache/spark/shuffle/local/pipelined/PipelinedChannelShuffleSuite.scala:
##########
@@ -685,6 +686,64 @@ 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 = consumer.partitions.indices.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 during setup.
+ blocks.take(cachedPartitions).foreach { block =>
+ assert(blockManager.putIterator(
+ block, Iterator((block.splitIndex, block.splitIndex)),
+ StorageLevel.MEMORY_ONLY, tellMaster = true))
+ }
+
assert(blockManager.master.getLocations(blocks.toArray[BlockId]).map(_.nonEmpty).toSeq
===
+ blocks.indices.map(_ < cachedPartitions))
+
+ val events = new JobEvents
+ sc.addSparkListener(events)
+ val error = intercept[SparkException](consumer.collect())
+ checkError(
+ exception = error,
+ condition = "PIPELINED_SHUFFLE_UNSUPPORTED",
+ sqlState = Some("0A000"),
+ parameters = Map("reason" ->
+ (s"cache/persist/localCheckpoint on RDD $consumer in a pipelined
consumer's " +
+ "within-stage chain (unpersist this RDD or rebuild it without
localCheckpoint " +
+ "before submitting the job; persist only RDDs that do not read a
pipelined " +
+ "shuffle within their stage)")))
+ sc.listenerBus.waitUntilEmpty()
+ assert(events.taskStarts.isEmpty)
+ assert(events.stageEnds.isEmpty)
+ assert(ChannelShuffleRendezvous.numQueuesForTesting === 0)
+ assert(consumer.getStorageLevel === StorageLevel.MEMORY_ONLY)
+
+ consumer.unpersist(blocking = true)
+ for (_ <- 0 until 2) {
+ assert(consumer.collect().map(_._1).sorted === (0 until
1000).toArray)
+ assert(ChannelShuffleRendezvous.numQueuesForTesting === 0)
+ }
+ }
+ }
+ }
+
+ test("cached producer input survives eviction between pipelined actions") {
Review Comment:
**nit: the name reads as if the cached blocks survive eviction, but the test
evicts them.** It removes the input's blocks with `removeRdd` and checks that
the second pipelined `collect()` recomputes the input and still succeeds. Could
the name say that, e.g. "pipelined actions recompute an evicted producer-input
cache"?
--
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]
