otterc commented on a change in pull request #30691: URL: https://github.com/apache/spark/pull/30691#discussion_r626226635
########## File path: core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala ########## @@ -3393,6 +3406,271 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assert(rprofsE === Set()) } + private def initPushBasedShuffleConfs(conf: SparkConf) = { + conf.set(config.SHUFFLE_SERVICE_ENABLED, true) + conf.set(config.PUSH_BASED_SHUFFLE_ENABLED, true) + conf.set("spark.master", "pushbasedshuffleclustermanager") + } + + test("SPARK-32920: shuffle merge finalization") { + initPushBasedShuffleConfs(conf) + DAGSchedulerSuite.clearMergerLocs + DAGSchedulerSuite.addMergerLocs(Seq("host1", "host2", "host3", "host4", "host5")) + val parts = 2 + val shuffleMapRdd = new MyRDD(sc, parts, Nil) + val shuffleDep = new ShuffleDependency(shuffleMapRdd, new HashPartitioner(parts)) + val reduceRdd = new MyRDD(sc, parts, List(shuffleDep), tracker = mapOutputTracker) + + // Submit a reduce job that depends which will create a map stage + submit(reduceRdd, (0 until parts).toArray) + completeShuffleMapStageSuccessfully(0, 0, parts) + assert(mapOutputTracker.getNumAvailableMergeResults(shuffleDep.shuffleId) == parts) + completeNextResultStageWithSuccess(1, 0) + assert(results === Map(0 -> 42, 1 -> 42)) + results.clear() + assertDataStructuresEmpty() + } + + test("SPARK-32920: merger locations not empty") { + initPushBasedShuffleConfs(conf) + conf.set(config.SHUFFLE_MERGER_LOCATIONS_MIN_STATIC_THRESHOLD, 3) + DAGSchedulerSuite.clearMergerLocs + DAGSchedulerSuite.addMergerLocs(Seq("host1", "host2", "host3", "host4", "host5")) + val parts = 2 + + val shuffleMapRdd = new MyRDD(sc, parts, Nil) + val shuffleDep = new ShuffleDependency(shuffleMapRdd, new HashPartitioner(parts)) + val reduceRdd = new MyRDD(sc, parts, List(shuffleDep), tracker = mapOutputTracker) + + // Submit a reduce job that depends which will create a map stage + submit(reduceRdd, (0 until parts).toArray) + completeShuffleMapStageSuccessfully(0, 0, parts) + val shuffleStage = scheduler.stageIdToStage(0).asInstanceOf[ShuffleMapStage] + assert(shuffleStage.shuffleDep.getMergerLocs.nonEmpty) + + assert(mapOutputTracker.getNumAvailableMergeResults(shuffleDep.shuffleId) == parts) + completeNextResultStageWithSuccess(1, 0) + assert(results === Map(0 -> 42, 1 -> 42)) + + results.clear() + assertDataStructuresEmpty() + } + + test("SPARK-32920: merger locations reuse from shuffle dependency") { + initPushBasedShuffleConfs(conf) + conf.set(config.SHUFFLE_MERGER_MAX_RETAINED_LOCATIONS, 3) + DAGSchedulerSuite.clearMergerLocs + DAGSchedulerSuite.addMergerLocs(Seq("host1", "host2", "host3", "host4", "host5")) + val parts = 2 + + val shuffleMapRdd = new MyRDD(sc, parts, Nil) + val shuffleDep = new ShuffleDependency(shuffleMapRdd, new HashPartitioner(parts)) + val reduceRdd = new MyRDD(sc, parts, List(shuffleDep), tracker = mapOutputTracker) + submit(reduceRdd, Array(0, 1)) + + completeShuffleMapStageSuccessfully(0, 0, parts) + assert(shuffleDep.getMergerLocs.nonEmpty) + val mergerLocs = shuffleDep.getMergerLocs + completeNextResultStageWithSuccess(1, 0 ) + + // submit another job w/ the shared dependency, and have a fetch failure + val reduce2 = new MyRDD(sc, 2, List(shuffleDep)) + submit(reduce2, Array(0, 1)) + // Note that the stage numbering here is only b/c the shared dependency produces a new, skipped + // stage. If instead it reused the existing stage, then this would be stage 2 + completeNextStageWithFetchFailure(3, 0, shuffleDep) + scheduler.resubmitFailedStages() + + assert(scheduler.runningStages.nonEmpty) + assert(scheduler.stageIdToStage(2) + .asInstanceOf[ShuffleMapStage].shuffleDep.getMergerLocs.nonEmpty) + val newMergerLocs = scheduler.stageIdToStage(2) + .asInstanceOf[ShuffleMapStage].shuffleDep.getMergerLocs + + // Check if same merger locs is reused for the new stage with shared shuffle dependency + assert(mergerLocs.zip(newMergerLocs).forall(x => x._1.host == x._2.host)) + completeShuffleMapStageSuccessfully(2, 0, 2) + completeNextResultStageWithSuccess(3, 1, idx => idx + 1234) + assert(results === Map(0 -> 1234, 1 -> 1235)) + + assertDataStructuresEmpty() + } + + test("SPARK-32920: Disable shuffle merge due to not enough mergers available") { + initPushBasedShuffleConfs(conf) + conf.set(config.SHUFFLE_MERGER_LOCATIONS_MIN_STATIC_THRESHOLD, 6) + DAGSchedulerSuite.clearMergerLocs + DAGSchedulerSuite.addMergerLocs(Seq("host1", "host2", "host3", "host4", "host5")) + val parts = 7 + + val shuffleMapRdd = new MyRDD(sc, parts, Nil) + val shuffleDep = new ShuffleDependency(shuffleMapRdd, new HashPartitioner(parts)) + val reduceRdd = new MyRDD(sc, parts, List(shuffleDep), tracker = mapOutputTracker) + + // Submit a reduce job that depends which will create a map stage + submit(reduceRdd, (0 until parts).toArray) + completeShuffleMapStageSuccessfully(0, 0, parts) + val shuffleStage = scheduler.stageIdToStage(0).asInstanceOf[ShuffleMapStage] + assert(shuffleStage.shuffleDep.mergerLocs.isEmpty) + + completeNextResultStageWithSuccess(1, 0) + assert(results === Map(2 -> 42, 5 -> 42, 4 -> 42, 1 -> 42, 3 -> 42, 6 -> 42, 0 -> 42)) + + results.clear() + assertDataStructuresEmpty() + } + + test("SPARK-32920: Ensure child stage should not start before all the" + + " parent stages are completed with shuffle merge finalized for all the parent stages") { Review comment: There are tests in `DAGSchedulerSuite` which have 4 spaces like: ``` test("[SPARK-19263] DAGScheduler should not submit multiple active tasksets," + " even with late completions from earlier stage attempts") ``` -- 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. For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org