jerrypeng commented on code in PR #59523:
URL: https://github.com/apache/spark/pull/59523#discussion_r4429459388
##########
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:
Renamed the test to “pipelined actions recompute an evicted producer-input
cache”. The cached blocks are removed while persistence intent remains set, and
the second pipelined action succeeds by recomputing the producer input.
--
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]
