[WIP] Stack/pipelined shuffle pr7 rtm

common/utils/src/main/resources/error/error-conditions.json

@@ -6399,6 +6399,18 @@
     ],
     "sqlState" : "42K03"
   },
+  "PIPELINED_SHUFFLE_CROSS_JOB_REUSE" : {
+    "message" : [
+      "The pipelined shuffle <shuffleId> is being reused across jobs. A pipelined (incrementally-readable) shuffle is transient and has no retained output for another job to read, so its producer stage cannot be shared across jobs."
+    ],
+    "sqlState" : "0A000"
+  },
+  "PIPELINED_SHUFFLE_UNSUPPORTED" : {
+    "message" : [
+      "A pipelined shuffle stage group cannot be scheduled because it uses an unsupported feature: <reason>. Pipelined (incrementally-readable) shuffles run their producer and consumer stages concurrently over a transient, once-through stream, which is incompatible with this feature in this version."
+    ],
+    "sqlState" : "0A000"
+  },
   "PIPELINE_DATASET_WITHOUT_FLOW" : {
     "message" : [
       "Pipeline dataset <identifier> does not have any defined flows. Please attach a query with the dataset's definition, or explicitly define at least one flow that writes to the dataset."

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,14 @@ 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 state is driver-only (the worker tracker caches nothing), so
+        // cleanup is a direct driver-side unregister. No RemoveShuffle RPC is needed: a streaming
+        // shuffle has no durable, block-manager-served files to remove.
+        logDebug("Cleaning pipelined shuffle " + shuffleId)
+        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 +320,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/StreamingShuffleOutputTracker.scala

@@ -230,6 +230,10 @@ private[spark] class StreamingShuffleOutputTrackerMaster(conf: SparkConf)
     }
   }
 
+  // Whether the given shuffle is registered. ContextCleaner uses this on the driver to skip
+  // regular (non-pipelined) shuffles, which are never registered here, before unregistering.
+  def containsShuffle(shuffleId: Int): Boolean = shuffleInfos.containsKey(shuffleId)
+
   // for testing purposes
   private[spark] def getShuffleInfo(shuffleId: Int): Option[StreamingShuffleInfo] = {
     Option(shuffleInfos.get(shuffleId))

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

@@ -63,4 +63,26 @@ private[spark] class ActiveJob(
   val finished = Array.fill[Boolean](numPartitions)(false)
 
   var numFinished = 0
+
+  /**
+   * Whether this job's RDD graph uses a `PipelinedShuffleDependency` (set once at job submission).
+   * Every pipelined-group scheduling path -- co-scheduling, deferral, and the per-submit
+   * `TaskSet.isPipelined` tagging -- is inert for a job without one, so this flag lets those paths
+   * short-circuit the group-membership graph walk for the common regular job at no cost.
+   */
+  var hasPipelinedDependency: Boolean = false
+
+  /**
+   * True once this job's pipelined group has passed up-front gang admission in
+   * `handleJobSubmitted` (or the slot check was disabled by config). This is a distinct fact from
+   * `hasPipelinedDependency` (which only says the job uses pipelining): the co-schedule path in
+   * `DAGScheduler.submitStage` gang-schedules a group's members with NO slot check, and checks this
+   * flag so that trust in up-front admission is enforced rather than merely commented.
+   *
+   * The whole job is a single pipelined group, so admission is a job-level fact. If a future change
+   * allows multiple groups per job, this job-level flag must be replaced by per-group admission
+   * state, and the check in `submitStage` replaced by that per-group gate -- NOT simply deleted, or
+   * an unadmitted group could be gang-scheduled and deadlock.
+   */
+  var pipelinedGroupAdmitted: Boolean = false
 }