[HADP-54419] Support ignore shuffle data migrated when decommission executor (apache#484)

Author: wangyum

core/src/main/scala/org/apache/spark/internal/config/package.scala

@@ -2438,6 +2438,14 @@ package object config {
       .stringConf
       .createWithDefaultString("PWR")
 
+  private[spark] val EXECUTOR_DECOMMISSION_IGNORE_SHUFFLE_DATA_MIGRATED =
+    ConfigBuilder("spark.executor.decommission.ignoreShuffleDataMigrated")
+      .doc("This configuration determines whether should ignore the shuffle data that " +
+        "has been migrated when it starts the decommissioning process.")
+      .version("3.5.0")
+      .booleanConf
+      .createWithDefault(true)
+
   private[spark] val STAGING_DIR = ConfigBuilder("spark.yarn.stagingDir")
     .doc("Staging directory used while submitting applications.")
     .version("2.0.0")

core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala

@@ -1180,9 +1180,15 @@ private[spark] class TaskSetManager(
     // could serve the shuffle outputs or the executor lost is caused by decommission (which
     // can destroy the whole host). The reason is the next stage wouldn't be able to fetch the
     // data from this dead executor so we would need to rerun these tasks on other executors.
-    val maybeShuffleMapOutputLoss = isShuffleMapTasks &&
-      !sched.sc.shuffleDriverComponents.supportsReliableStorage() &&
-      (reason.isInstanceOf[ExecutorDecommission] || !env.blockManager.externalShuffleServiceEnabled)
+    val maybeShuffleMapOutputLoss =
+      if (conf.get(config.EXECUTOR_DECOMMISSION_IGNORE_SHUFFLE_DATA_MIGRATED)) {
+        isShuffleMapTasks && !sched.sc.shuffleDriverComponents.supportsReliableStorage() &&
+          !env.blockManager.externalShuffleServiceEnabled
+      } else {
+        isShuffleMapTasks && !sched.sc.shuffleDriverComponents.supportsReliableStorage() &&
+          (reason.isInstanceOf[ExecutorDecommission] ||
+            !env.blockManager.externalShuffleServiceEnabled)
+      }
     if (maybeShuffleMapOutputLoss && !isZombie) {
       for ((tid, info) <- taskInfos if info.executorId == execId) {
         val index = info.index

core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala

@@ -861,7 +861,8 @@ class TaskSetManagerSuite
   }
 
   test("SPARK-32653: Decommissioned host should not be used to calculate locality levels") {
-    sc = new SparkContext("local", "test")
+    val conf = new SparkConf().set(config.EXECUTOR_DECOMMISSION_IGNORE_SHUFFLE_DATA_MIGRATED, false)
+    sc = new SparkContext("local", "test", conf)
     sched = new FakeTaskScheduler(sc)
     val backend = mock(classOf[SchedulerBackend])
     doNothing().when(backend).reviveOffers()