jerrypeng commented on code in PR #59523:
URL: https://github.com/apache/spark/pull/59523#discussion_r4408514165
##########
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:
Agreed. I added "cache" to the existing shared-ancestry test matrix,
covering both one and 32 cached roots, and now assert exactly five reads per
shared edge. This catches an additional traversal that the previous <= 6
assertion would allow. I removed the separate cache test and its duplicated
setup.
--
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]
