[SPARK-59138][CORE] Preserve reliably-stored shuffle map outputs on executor loss

core/src/main/scala/org/apache/spark/MapOutputTracker.scala

@@ -63,7 +63,8 @@ import org.apache.spark.util.io.{ChunkedByteBuffer, ChunkedByteBufferOutputStrea
 private class ShuffleStatus(
     numPartitions: Int,
     numReducers: Int = -1,
-    bufferRacingMigrations: Boolean = false) extends Logging {
+    bufferRacingMigrations: Boolean = false,
+    val isReliablyStored: Boolean = false) extends Logging {
 
   private val (readLock, writeLock) = {
     val lock = new ReentrantReadWriteLock()
@@ -929,23 +930,37 @@ private[spark] class MapOutputTrackerMaster(
     shuffleStatuses.valuesIterator.count(_.hasCachedSerializedBroadcast)
   }
 
-  def registerShuffle(shuffleId: Int, numMaps: Int, numReduces: Int): Unit = {
+  def registerShuffle(shuffleId: Int, numMaps: Int, numReduces: Int): Unit =
+    registerShuffle(shuffleId, numMaps, numReduces, isReliablyStored = false)
+
+  def registerShuffle(
+      shuffleId: Int,
+      numMaps: Int,
+      numReduces: Int,
+      isReliablyStored: Boolean): Unit = {
     if (pushBasedShuffleEnabled) {
       if (shuffleStatuses.put(shuffleId,
-        new ShuffleStatus(numMaps, numReduces, bufferRacingMigrations)).isDefined) {
+        new ShuffleStatus(numMaps, numReduces, bufferRacingMigrations,
+          isReliablyStored)).isDefined) {
         throw new IllegalArgumentException("Shuffle ID " + shuffleId + " registered twice")
       }
     } else {
       if (shuffleStatuses.put(shuffleId,
-        new ShuffleStatus(numMaps, bufferRacingMigrations = bufferRacingMigrations)).isDefined) {
+        new ShuffleStatus(numMaps, bufferRacingMigrations = bufferRacingMigrations,
+          isReliablyStored = isReliablyStored)).isDefined) {
         throw new IllegalArgumentException("Shuffle ID " + shuffleId + " registered twice")
       }
     }
   }
 
   // ShuffleOutputTrackerMaster: a regular shuffle has no per-job registration, so jobId is ignored.
-  override def registerShuffle(shuffleId: Int, numMaps: Int, numReduces: Int, jobId: Int): Unit =
-    registerShuffle(shuffleId, numMaps, numReduces)
+  override def registerShuffle(
+      shuffleId: Int,
+      numMaps: Int,
+      numReduces: Int,
+      jobId: Int,
+      isReliablyStored: Boolean): Unit =
+    registerShuffle(shuffleId, numMaps, numReduces, isReliablyStored)
 
   def updateMapOutput(shuffleId: Int, mapId: Long, bmAddress: BlockManagerId): Unit = {
     shuffleStatuses.get(shuffleId) match {
@@ -1044,20 +1059,48 @@ private[spark] class MapOutputTrackerMaster(
   /**
    * Removes all shuffle outputs associated with this host. Note that this will also remove
    * outputs which are served by an external shuffle server (if one exists).
+   *
+   * When `skipReliablyStored` is true (executor/worker loss rather than a fetch failure),
+   * shuffles whose output is reliably stored off-executor are left intact, since losing the host
+   * does not lose their output.
    */
-  def removeOutputsOnHost(host: String): Unit = {
-    shuffleStatuses.valuesIterator.foreach { _.removeOutputsOnHost(host) }
-    incrementEpoch()
+  def removeOutputsOnHost(host: String): Unit =
+    removeOutputsOnHost(host, skipReliablyStored = false)
+
+  def removeOutputsOnHost(host: String, skipReliablyStored: Boolean): Unit = {
+    var removedAny = false
+    shuffleStatuses.valuesIterator.foreach { status =>
+      if (!(skipReliablyStored && status.isReliablyStored)) {
+        status.removeOutputsOnHost(host)
+        removedAny = true
+      }
+    }
+    // Skip the epoch bump when nothing was removed (every shuffle was reliably stored): a bump
+    // needlessly invalidates every executor's cached map statuses and forces a re-fetch.
+    if (removedAny) incrementEpoch()
   }
 
   /**
    * Removes all shuffle outputs associated with this executor. Note that this will also remove
    * outputs which are served by an external shuffle server (if one exists), as they are still
    * registered with this execId.
+   *
+   * When `skipReliablyStored` is true (executor loss rather than a fetch failure), shuffles whose
+   * output is reliably stored off-executor are left intact: losing the executor does not lose
+   * their output, so unregistering would force a needless map-stage recompute.
    */
-  def removeOutputsOnExecutor(execId: String): Unit = {
-    shuffleStatuses.valuesIterator.foreach { _.removeOutputsOnExecutor(execId) }
-    incrementEpoch()
+  def removeOutputsOnExecutor(execId: String): Unit =
+    removeOutputsOnExecutor(execId, skipReliablyStored = false)
+
+  def removeOutputsOnExecutor(execId: String, skipReliablyStored: Boolean): Unit = {
+    var removedAny = false
+    shuffleStatuses.valuesIterator.foreach { status =>
+      if (!(skipReliablyStored && status.isReliablyStored)) {
+        status.removeOutputsOnExecutor(execId)
+        removedAny = true
+      }
+    }
+    if (removedAny) incrementEpoch()
   }
 
   /**
@@ -1083,6 +1126,14 @@ private[spark] class MapOutputTrackerMaster(
   /** Check if the given shuffle is being tracked */
   override def containsShuffle(shuffleId: Int): Boolean = shuffleStatuses.contains(shuffleId)
 
+  /**
+   * Whether this shuffle's output is reliably stored off-executor, so it is not lost when an
+   * executor or worker holding it is lost. Unknown shuffles default to false.
+   */
+  def isReliablyStored(shuffleId: Int): Boolean = {
+    shuffleStatuses.get(shuffleId).exists(_.isReliablyStored)
+  }
+
   def getNumAvailableOutputs(shuffleId: Int): Int = {
     shuffleStatuses.get(shuffleId).map(_.numAvailableMapOutputs).getOrElse(0)
   }

core/src/main/scala/org/apache/spark/StreamingShuffleOutputTracker.scala

@@ -38,7 +38,12 @@ import org.apache.spark.util.ThreadUtils
  */
 private[spark] trait ShuffleOutputTrackerMaster {
   /** Register a shuffle so its outputs can be tracked. `jobId` is used by the streaming tracker. */
-  def registerShuffle(shuffleId: Int, numMaps: Int, numReduces: Int, jobId: Int): Unit
+  def registerShuffle(
+      shuffleId: Int,
+      numMaps: Int,
+      numReduces: Int,
+      jobId: Int,
+      isReliablyStored: Boolean = false): Unit
   /** Whether the given shuffle is registered with this tracker. */
   def containsShuffle(shuffleId: Int): Boolean
   /** Unregister a shuffle and release its tracked state. */
@@ -237,7 +242,12 @@ private[spark] class StreamingShuffleOutputTrackerMaster(conf: SparkConf)
     pool
   }
 
-  override def registerShuffle(shuffleId: Int, numMaps: Int, numReduces: Int, jobId: Int): Unit = {
+  override def registerShuffle(
+      shuffleId: Int,
+      numMaps: Int,
+      numReduces: Int,
+      jobId: Int,
+      isReliablyStored: Boolean): Unit = {
     logInfo(log"Registering shuffleId ${MDC(LogKeys.SHUFFLE_ID, shuffleId)} with ${
       MDC(LogKeys.NUM_MAPPERS, numMaps)} mappers and ${
       MDC(LogKeys.NUM_REDUCERS, numReduces)} reducers")

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

@@ -762,7 +762,8 @@ private[spark] class DAGScheduler(
         log"(${MDC(CREATION_SITE, rdd.getCreationSite)}) as input to " +
         log"shuffle ${MDC(SHUFFLE_ID, shuffleDep.shuffleId)}")
       outputTracker.registerShuffle(shuffleDep.shuffleId, rdd.partitions.length,
-        shuffleDep.partitioner.numPartitions, jobId)
+        shuffleDep.partitioner.numPartitions, jobId,
+        isReliablyStored = shuffleDep.shuffleHandle.isReliablyStored)
     }
     stage
   }
@@ -4199,7 +4200,10 @@ private[spark] class DAGScheduler(
       execId = execId,
       fileLost = fileLost,
       hostToUnregisterOutputs = workerHost,
-      maybeEpoch = None)
+      maybeEpoch = None,
+      // Executor loss (not a fetch failure): preserve shuffles whose output is reliably stored
+      // off-executor. Their data survives the executor, so recomputing them would be wasteful.
+      skipReliablyStored = true)
   }
 
   /**
@@ -4267,7 +4271,8 @@ private[spark] class DAGScheduler(
       fileLost: Boolean,
       hostToUnregisterOutputs: Option[String],
       maybeEpoch: Option[Long] = None,
-      ignoreShuffleFileLostEpoch: Boolean = false): Unit = {
+      ignoreShuffleFileLostEpoch: Boolean = false,
+      skipReliablyStored: Boolean = false): Unit = {
     val currentEpoch = maybeEpoch.getOrElse(mapOutputTracker.getEpoch)
     logDebug(s"Considering removal of executor $execId; " +
       s"fileLost: $fileLost, currentEpoch: $currentEpoch")
@@ -4307,11 +4312,11 @@ private[spark] class DAGScheduler(
           case Some(host) =>
             logInfo(log"Shuffle files lost for host: ${MDC(HOST, host)} (epoch " +
               log"${MDC(EPOCH, currentEpoch)}")
-            mapOutputTracker.removeOutputsOnHost(host)
+            mapOutputTracker.removeOutputsOnHost(host, skipReliablyStored)
           case None =>
               logInfo(log"Shuffle files lost for executor: ${MDC(EXECUTOR_ID, execId)} " +
                 log"(epoch ${MDC(EPOCH, currentEpoch)})")
-            mapOutputTracker.removeOutputsOnExecutor(execId)
+            mapOutputTracker.removeOutputsOnExecutor(execId, skipReliablyStored)
         }
       }
     }
@@ -4334,7 +4339,9 @@ private[spark] class DAGScheduler(
       message: String): Unit = {
     logInfo(log"Shuffle files lost for worker ${MDC(WORKER_ID, workerId)} " +
       log"on host ${MDC(HOST, host)}")
-    mapOutputTracker.removeOutputsOnHost(host)
+    // Worker loss (not a fetch failure): reliably-stored shuffle output lives off the worker and
+    // survives, so leave those outputs registered and only drop the rest.
+    mapOutputTracker.removeOutputsOnHost(host, skipReliablyStored = true)
     clearCacheLocs()
   }
 

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

@@ -1228,8 +1228,20 @@ private[spark] class TaskSetManager(
     // pipelined set and aborts the whole group. Note isZombie already skips a fully-complete
     // producer's set; this guard also covers a PARTIALLY-complete producer losing an executor on
     // decommission.
+
+    // OR, not AND: a shuffle reliably stored off-executor (globally, or just this one via a remote
+    // shuffle service) keeps its map output when the executor dies. The per-shuffle bit only ever
+    // adds reliability (defaults to false, set true solely by an opting-in manager), so a false
+    // there means "no info", not "unreliable".
+    val reliablyStored = sched.sc.shuffleDriverComponents.supportsReliableStorage() ||
+      taskSet.shuffleId.exists { shuffleId =>
+        sched.mapOutputTracker match {
+          case master: MapOutputTrackerMaster => master.isReliablyStored(shuffleId)
+          case _ => false
+        }
+      }
     val maybeShuffleMapOutputLoss = isShuffleMapTasks && !taskSet.isPipelined &&
-      !sched.sc.shuffleDriverComponents.supportsReliableStorage() &&
+      !reliablyStored &&
       (reason.isInstanceOf[ExecutorDecommission] || !env.blockManager.externalShuffleServiceEnabled)
     if (maybeShuffleMapOutputLoss && !isZombie) {
       val iter1 = taskIdsOnExec.iterator

core/src/main/scala/org/apache/spark/shuffle/ShuffleHandle.scala

@@ -25,4 +25,12 @@ import org.apache.spark.annotation.DeveloperApi
  * @param shuffleId ID of the shuffle
  */
 @DeveloperApi
-abstract class ShuffleHandle(val shuffleId: Int) extends Serializable {}
+abstract class ShuffleHandle(val shuffleId: Int) extends Serializable {
+  /**
+   * Whether this shuffle's output is stored reliably outside the executors that produced it (e.g.
+   * a remote shuffle service). When true, losing an executor does not lose this shuffle's output,
+   * so its map outputs are not unregistered on executor loss. Defaults to false; a ShuffleManager
+   * that routes a shuffle to reliable storage overrides this on the handle it returns.
+   */
+  def isReliablyStored: Boolean = false
+}

core/src/test/scala/org/apache/spark/MapOutputTrackerSuite.scala

@@ -136,6 +136,46 @@ class MapOutputTrackerSuite extends SparkFunSuite with LocalSparkContext {
     rpcEnv.shutdown()
   }
 
+  test("SPARK-59138: executor loss skips reliably-stored shuffles but not local-disk ones") {
+    val rpcEnv = createRpcEnv("test")
+    val tracker = newTrackerMaster()
+    tracker.trackerEndpoint = rpcEnv.setupEndpoint(MapOutputTracker.ENDPOINT_NAME,
+      new MapOutputTrackerMasterEndpoint(rpcEnv, tracker, conf))
+
+    val size = MapStatus.compressSize(1000L)
+    // Shuffle 0: local-disk (not reliably stored). Shuffle 1: reliably stored off-executor.
+    tracker.registerShuffle(0, 1, MergeStatus.SHUFFLE_PUSH_DUMMY_NUM_REDUCES)
+    tracker.registerShuffle(1, 1, MergeStatus.SHUFFLE_PUSH_DUMMY_NUM_REDUCES,
+      isReliablyStored = true)
+    tracker.registerMapOutput(0, 0, MapStatus(BlockManagerId("a", "hostA", 1000), Array(size), 5))
+    tracker.registerMapOutput(1, 0, MapStatus(BlockManagerId("a", "hostA", 1000), Array(size), 6))
+
+    assert(tracker.isReliablyStored(0) === false)
+    assert(tracker.isReliablyStored(1) === true)
+
+    // Executor loss: skip reliably-stored shuffles. Shuffle 0 drops, shuffle 1 stays.
+    tracker.removeOutputsOnExecutor("a", skipReliablyStored = true)
+    assert(tracker.getNumAvailableOutputs(0) === 0)
+    assert(tracker.getNumAvailableOutputs(1) === 1)
+
+    // Losing an executor when only reliably-stored shuffles remain removes nothing, so the
+    // epoch must not bump (a bump would needlessly invalidate every executor's cached statuses).
+    tracker.unregisterShuffle(0)
+    val epochBeforeNoOp = tracker.getEpoch
+    tracker.removeOutputsOnExecutor("a", skipReliablyStored = true)
+    assert(tracker.getEpoch === epochBeforeNoOp)
+
+    // Fetch failure (skip = false): even the reliably-stored shuffle's output is removed, and
+    // because something was removed the epoch bumps.
+    val epochBeforeRemoval = tracker.getEpoch
+    tracker.removeOutputsOnExecutor("a", skipReliablyStored = false)
+    assert(tracker.getNumAvailableOutputs(1) === 0)
+    assert(tracker.getEpoch > epochBeforeRemoval)
+
+    tracker.stop()
+    rpcEnv.shutdown()
+  }
+
   test("remote fetch") {
     val hostname = "localhost"
     val rpcEnv = createRpcEnv("spark", hostname, 0, new SecurityManager(conf))