jerrypeng commented on code in PR #59523:
URL: https://github.com/apache/spark/pull/59523#discussion_r4408574780
##########
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:
Agreed. I retained the persisted and local-checkpoint variants as
cache-preflight precedence tests and renamed them accordingly. The uncached
variant continues to test repeated consumption. I removed the misleading test
comment and updated the repeated-read Scaladoc to identify the separate cache
preflight that rejects consumer persistence.
##########
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:
Agreed. I removed the three added ordinary-shuffle cases. Existing
cache-pruning tests and the ordinary-job traversal test retain the relevant
coverage; these additional cases did not exercise the new persistence
restriction.
--
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]
