[WIP] Stack/pipelined shuffle pr8 rtm

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

@@ -236,6 +236,10 @@ private[spark] class ContextCleaner(
   /** Perform shuffle cleanup. */
   def doCleanupShuffle(shuffleId: Int, blocking: Boolean): Unit = {
     try {
+      // A shuffle lives in exactly one tracker, split by dependency type: a regular shuffle in the
+      // MapOutputTracker, a pipelined shuffle in the driver-only StreamingShuffleOutputTracker (see
+      // DAGScheduler.createShuffleMapStage). Clean up whichever holds it -- the two branches are
+      // independent, each keyed on its own tracker's membership, so neither depends on the other.
       if (mapOutputTrackerMaster.containsShuffle(shuffleId)) {
         logDebug("Cleaning shuffle " + shuffleId)
         // Shuffle must be removed before it's unregistered from the output tracker
@@ -244,6 +248,19 @@ private[spark] class ContextCleaner(
         mapOutputTrackerMaster.unregisterShuffle(shuffleId)
         listeners.asScala.foreach(_.shuffleCleaned(shuffleId))
         logDebug("Cleaned shuffle " + shuffleId)
+      } else if (streamingShuffleOutputTrackerMaster.exists(_.containsShuffle(shuffleId))) {
+        // A pipelined shuffle's output state is driver-only (the worker tracker caches nothing), so
+        // it is unregistered directly from the StreamingShuffleOutputTracker rather than the
+        // MapOutputTracker. Still call shuffleDriverComponents.removeShuffle to balance the
+        // registerShuffle that the ShuffleDependency constructor issues for EVERY shuffle (incl. a
+        // pipelined one): a custom components impl may track per-shuffle driver state that would
+        // otherwise leak. For the default impl this is just a RemoveShuffle RPC that finds no
+        // blocks (a pipelined shuffle has none), matching the regular branch's ordering.
+        logDebug("Cleaning pipelined shuffle " + shuffleId)
+        shuffleDriverComponents.removeShuffle(shuffleId, blocking)
+        streamingShuffleOutputTrackerMaster.foreach(_.unregisterShuffle(shuffleId))
+        listeners.asScala.foreach(_.shuffleCleaned(shuffleId))
+        logDebug("Cleaned pipelined shuffle " + shuffleId)
       } else {
         logDebug("Asked to cleanup non-existent shuffle (maybe it was already removed)")
       }
@@ -308,6 +325,8 @@ private[spark] class ContextCleaner(
 
   private def broadcastManager = sc.env.broadcastManager
   private def mapOutputTrackerMaster = sc.env.mapOutputTracker.asInstanceOf[MapOutputTrackerMaster]
+  private def streamingShuffleOutputTrackerMaster: Option[StreamingShuffleOutputTrackerMaster] =
+    sc.env.streamingShuffleOutputTracker.map(_.asInstanceOf[StreamingShuffleOutputTrackerMaster])
 }
 
 private object ContextCleaner {

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

@@ -799,7 +799,7 @@ private[spark] class MapOutputTrackerMaster(
     conf: SparkConf,
     private[spark] val broadcastManager: BroadcastManager,
     private[spark] val isLocal: Boolean)
-  extends MapOutputTracker(conf) {
+  extends MapOutputTracker(conf) with ShuffleOutputTrackerMaster {
 
   // The size at which we use Broadcast to send the map output statuses to the executors
   private val minSizeForBroadcast = conf.get(SHUFFLE_MAPOUTPUT_MIN_SIZE_FOR_BROADCAST).toInt
@@ -943,6 +943,10 @@ private[spark] class MapOutputTrackerMaster(
     }
   }
 
+  // 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)
+
   def updateMapOutput(shuffleId: Int, mapId: Long, bmAddress: BlockManagerId): Unit = {
     shuffleStatuses.get(shuffleId) match {
       case Some(shuffleStatus) =>
@@ -1030,7 +1034,7 @@ private[spark] class MapOutputTrackerMaster(
   }
 
   /** Unregister shuffle data */
-  def unregisterShuffle(shuffleId: Int): Unit = {
+  override def unregisterShuffle(shuffleId: Int): Unit = {
     shuffleStatuses.remove(shuffleId).foreach { shuffleStatus =>
       shuffleStatus.invalidateSerializedMapOutputStatusCache()
       shuffleStatus.invalidateSerializedMergeOutputStatusCache()
@@ -1077,7 +1081,7 @@ private[spark] class MapOutputTrackerMaster(
   }
 
   /** Check if the given shuffle is being tracked */
-  def containsShuffle(shuffleId: Int): Boolean = shuffleStatuses.contains(shuffleId)
+  override def containsShuffle(shuffleId: Int): Boolean = shuffleStatuses.contains(shuffleId)
 
   def getNumAvailableOutputs(shuffleId: Int): Int = {
     shuffleStatuses.get(shuffleId).map(_.numAvailableMapOutputs).getOrElse(0)

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

@@ -28,6 +28,23 @@ import org.apache.spark.internal.config.SHUFFLE_MAPOUTPUT_DISPATCHER_NUM_THREADS
 import org.apache.spark.rpc.{RpcCallContext, RpcEndpoint, RpcEndpointRef, RpcEnv}
 import org.apache.spark.util.ThreadUtils
 
+/**
+ * The driver-side registry for a shuffle's output. A regular shuffle is served by the
+ * `MapOutputTrackerMaster`, a pipelined (streaming) shuffle by the
+ * `StreamingShuffleOutputTrackerMaster` -- split by dependency type, with no overlap. This is the
+ * small common surface the DAGScheduler and ContextCleaner drive polymorphically, so they select
+ * the right tracker by shuffle-dependency type (see `DAGScheduler.outputTrackerMaster`) rather than
+ * special-casing pipelined shuffles at each call site.
+ */
+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
+  /** Whether the given shuffle is registered with this tracker. */
+  def containsShuffle(shuffleId: Int): Boolean
+  /** Unregister a shuffle and release its tracked state. */
+  def unregisterShuffle(shuffleId: Int): Unit
+}
+
 private[spark] sealed trait StreamingShuffleTaskLocationTrackerMessage
 
 private[spark] case class UpdateStreamingShuffleTaskLocation(
@@ -196,7 +213,7 @@ private[spark] abstract class StreamingShuffleOutputTracker(conf: SparkConf) ext
 private[spark] case class StreamingShuffleInfo(numMaps: Int, numReduces: Int, jobId: Int)
 
 private[spark] class StreamingShuffleOutputTrackerMaster(conf: SparkConf)
-  extends StreamingShuffleOutputTracker(conf) {
+  extends StreamingShuffleOutputTracker(conf) with ShuffleOutputTrackerMaster {
 
   // map that stores task location information organized in the following fashion
   // shuffle id -> {mapId -> location}
@@ -220,7 +237,7 @@ private[spark] class StreamingShuffleOutputTrackerMaster(conf: SparkConf)
     pool
   }
 
-  def registerShuffle(shuffleId: Int, numMaps: Int, numReduces: Int, jobId: Int): Unit = {
+  override def registerShuffle(shuffleId: Int, numMaps: Int, numReduces: Int, jobId: Int): Unit = {
     logInfo(log"Registering shuffleId ${MDC(LogKeys.SHUFFLE_ID, shuffleId)} with ${
       MDC(LogKeys.NUM_MAPPERS, numMaps)} mappers and ${
       MDC(LogKeys.NUM_REDUCERS, numReduces)} reducers")
@@ -230,6 +247,8 @@ private[spark] class StreamingShuffleOutputTrackerMaster(conf: SparkConf)
     }
   }
 
+  override def containsShuffle(shuffleId: Int): Boolean = shuffleInfos.containsKey(shuffleId)
+
   // for testing purposes
   private[spark] def getShuffleInfo(shuffleId: Int): Option[StreamingShuffleInfo] = {
     Option(shuffleInfos.get(shuffleId))