diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index 7b31c6e253576..2e495416d67fd 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -6399,6 +6399,18 @@ ], "sqlState" : "42K03" }, + "PIPELINED_SHUFFLE_CROSS_JOB_REUSE" : { + "message" : [ + "The pipelined shuffle 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: . 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 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." diff --git a/core/src/main/scala/org/apache/spark/ContextCleaner.scala b/core/src/main/scala/org/apache/spark/ContextCleaner.scala index 0b3c22a22cb46..a01a480147d72 100644 --- a/core/src/main/scala/org/apache/spark/ContextCleaner.scala +++ b/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 { diff --git a/core/src/main/scala/org/apache/spark/StreamingShuffleOutputTracker.scala b/core/src/main/scala/org/apache/spark/StreamingShuffleOutputTracker.scala index 090be7f691d4b..9d6fc4f1936c8 100644 --- a/core/src/main/scala/org/apache/spark/StreamingShuffleOutputTracker.scala +++ b/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)) diff --git a/core/src/main/scala/org/apache/spark/scheduler/ActiveJob.scala b/core/src/main/scala/org/apache/spark/scheduler/ActiveJob.scala index 9876668194a84..4491461ee1199 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/ActiveJob.scala +++ b/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 } diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index bfd24ceb2f895..7b07f62ade832 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -45,7 +45,7 @@ import org.apache.spark.network.shuffle.{BlockStoreClient, MergeFinalizerListene import org.apache.spark.network.shuffle.protocol.MergeStatuses import org.apache.spark.network.util.JavaUtils import org.apache.spark.partial.{ApproximateActionListener, ApproximateEvaluator, PartialResult} -import org.apache.spark.rdd.{RDD, RDDCheckpointData} +import org.apache.spark.rdd.{DeterministicLevel, RDD, RDDCheckpointData, ReliableRDDCheckpointData} import org.apache.spark.resource.{CpuAmount, ResourceProfile, TaskResourceProfile} import org.apache.spark.resource.ResourceProfile.{CPUS, DEFAULT_RESOURCE_PROFILE_ID, EXECUTOR_CORES_LOCAL_PROPERTY, MAX_TASKS_PER_EXECUTOR_LOCAL_PROPERTY, PYSPARK_MEMORY_LOCAL_PROPERTY} import org.apache.spark.rpc.RpcTimeout @@ -187,6 +187,11 @@ private[spark] class DAGScheduler( delayedTaskCompletionEvents: ListBuffer[CompletionEvent] = new ListBuffer[CompletionEvent]) private[scheduler] val dependentStageMap = new HashMap[Stage, DependentStageInfo] + // Whether we have already logged that the pipelined-group slot check is disabled. Only the + // (single-threaded) event loop touches this, so a plain var is safe. Keeps the warning to once + // per scheduler rather than once per submitted batch job. + private var warnedPipelinedSlotCheckDisabled = false + private[scheduler] val activeJobs = new HashSet[ActiveJob] // Track all the jobs submitted by the same query execution, will clean up after @@ -622,6 +627,21 @@ private[spark] class DAGScheduler( firstJobId: Int): ShuffleMapStage = { shuffleIdToMapStage.get(shuffleDep.shuffleId) match { case Some(stage) => + // A pipelined shuffle is transient: it is a once-through live stream with no retained, + // addressable output, so reusing its producer stage across jobs is unsound (there is no + // durable output for a second job to read). Reuse must be prevented explicitly -- from the + // scheduler's view a shuffle-map stage can be reused unless something forbids it. If a + // pipelined dependency's shuffleId is already bound to a stage from a different job, that + // is the forbidden cross-job reuse; fail fast. (Within the same job the cached stage is the + // one we just created, so returning it is correct and not reuse.) + if (shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]] && + !stage.jobIds.contains(firstJobId)) { + throw new SparkException( + errorClass = "PIPELINED_SHUFFLE_CROSS_JOB_REUSE", + messageParameters = scala.collection.immutable.Map( + "shuffleId" -> shuffleDep.shuffleId.toString), + cause = null) + } stage case None => @@ -671,6 +691,25 @@ private[spark] class DAGScheduler( checkBarrierStageWithDynamicAllocation(rdd) checkBarrierStageWithNumSlots(rdd, resourceProfile) checkBarrierStageWithRDDChainPattern(rdd, rdd.getNumPartitions) + checkPipelinedProducerSupported(shuffleDep) + // A pipelined shuffle is served solely by the StreamingShuffleOutputTracker, so that tracker + // MUST exist for one -- it is created whenever a streaming-capable shuffle manager is set up + // (see SparkEnv.initializeStreamingShuffleOutputTracker), a prerequisite for producing a + // PipelinedShuffleDependency. Resolve it up front (BEFORE any stage-map mutation below), so a + // fail-loud on the misconfiguration leaves no partial scheduler state -- and so a pipelined + // shuffle is never silently registered in no tracker (a consumer would then find no writer + // locations). The reader enforces the same invariant (see StreamingShuffleReader). None for a + // regular shuffle, which is registered with the MapOutputTracker instead. + val streamingTrackerForPipelined: Option[StreamingShuffleOutputTrackerMaster] = + if (shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) { + Some(sc.env.streamingShuffleOutputTracker + .getOrElse(throw new IllegalStateException( + s"A pipelined shuffle (id ${shuffleDep.shuffleId}) requires a " + + "StreamingShuffleOutputTracker, but none is configured")) + .asInstanceOf[StreamingShuffleOutputTrackerMaster]) + } else { + None + } val numTasks = rdd.partitions.length val parents = getOrCreateParentStages(shuffleDeps, jobId) val id = nextStageId.getAndIncrement() @@ -682,18 +721,100 @@ private[spark] class DAGScheduler( shuffleIdToMapStage(shuffleDep.shuffleId) = stage updateJobIdStageIdMaps(jobId, stage) - if (!mapOutputTracker.containsShuffle(shuffleDep.shuffleId)) { - // Kind of ugly: need to register RDDs with the cache and map output tracker here - // since we can't do it in the RDD constructor because # of partitions is unknown - logInfo(log"Registering RDD ${MDC(RDD_ID, rdd.id)} " + - log"(${MDC(CREATION_SITE, rdd.getCreationSite)}) as input to " + - log"shuffle ${MDC(SHUFFLE_ID, shuffleDep.shuffleId)}") - mapOutputTracker.registerShuffle(shuffleDep.shuffleId, rdd.partitions.length, - shuffleDep.partitioner.numPartitions) + // A pipelined shuffle and a regular shuffle live in DIFFERENT trackers -- the split is by + // dependency type, with no overlap. A pipelined shuffle produces no durable, addressable map + // output; its producer/consumer task-location routing lives entirely in the + // StreamingShuffleOutputTracker (resolved fail-loud above), and it is never registered with the + // MapOutputTracker (nothing reads such an entry -- availability is tracked on the stage via + // pipelinedCompletedPartitions, the reader locates writers through the streaming tracker, and + // every MapOutputTracker path a pipelined shuffle could reach is either type-gated out or + // tolerant of its absence). A regular shuffle registers with the MapOutputTracker as before. + streamingTrackerForPipelined match { + case Some(streamingTracker) => + // Register the producer/consumer fan-in before any consumer task asks for a writer loc. + // Self-guarded on the tracker's own state (createShuffleMapStage runs once per shuffleId + // via getOrCreateShuffleMapStage, but guard defensively against a re-entry). + if (!streamingTracker.containsShuffle(shuffleDep.shuffleId)) { + logInfo(log"Registering RDD ${MDC(RDD_ID, rdd.id)} " + + log"(${MDC(CREATION_SITE, rdd.getCreationSite)}) as input to pipelined " + + log"shuffle ${MDC(SHUFFLE_ID, shuffleDep.shuffleId)}") + streamingTracker.registerShuffle( + shuffleDep.shuffleId, + rdd.partitions.length, + shuffleDep.partitioner.numPartitions, + jobId) + } + case None if !mapOutputTracker.containsShuffle(shuffleDep.shuffleId) => + // Kind of ugly: need to register RDDs with the cache and map output tracker here + // since we can't do it in the RDD constructor because # of partitions is unknown + logInfo(log"Registering RDD ${MDC(RDD_ID, rdd.id)} " + + log"(${MDC(CREATION_SITE, rdd.getCreationSite)}) as input to " + + log"shuffle ${MDC(SHUFFLE_ID, shuffleDep.shuffleId)}") + mapOutputTracker.registerShuffle(shuffleDep.shuffleId, rdd.partitions.length, + shuffleDep.partitioner.numPartitions) + + case _ => // regular shuffle already registered with the MapOutputTracker; nothing to do } stage } + private def pipelinedUnsupportedError(reason: String): PipelinedShuffleUnsupportedException = + new PipelinedShuffleUnsupportedException(reason) + + /** + * Fail-fast on producer-side idioms a pipelined shuffle cannot support, checked when the producer + * stage is created. A pipelined shuffle runs its producer and consumer stages concurrently over a + * transient, once-through stream that a group never recomputes in isolation (any failure aborts + * the whole group), so mechanisms that recompute/roll back a single stage are moot, and features + * that expose output only after a global barrier are incompatible with + * incremental reads. Rejecting here (before the stage is used) keeps a misuse from silently + * mis-scheduling. Inert for a regular ShuffleDependency. + * + * Group-level idioms are handled elsewhere, since they are properties of the group rather than a + * single producer stage: fan-out (a producer with more than one consumer) and a group with a + * non-default resource profile are rejected up front at job submission by + * `checkPipelinedGroupsSupportedInRDDGraph` (before any stage is created). A regular shuffle + * internal to a group does not arise for the all-pipelined job shape (groups are split at + * regular-shuffle boundaries) and so is not checked. + */ + private def checkPipelinedProducerSupported(shuffleDep: ShuffleDependency[_, _, _]): Unit = { + if (!shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) { + return + } + val rdd: RDD[_] = shuffleDep.rdd + // Barrier: exposes output only after a global sync, contradicting concurrent partial reads. + if (rdd.isBarrier()) { + throw pipelinedUnsupportedError("barrier execution in a pipelined-group member stage") + } + // Dynamic resource allocation: gang admission needs a stable slot set; reclaiming executors + // from a pinned-open group can deadlock it. + if (Utils.isDynamicAllocationEnabled(sc.conf)) { + throw pipelinedUnsupportedError("dynamic resource allocation with a pipelined shuffle") + } + // Statically-indeterminate producer: its recovery is stage rollback-and-recompute, which a + // group never performs (any failure aborts the whole group); reject rather than carry dead + // machinery. + if (rdd.outputDeterministicLevel == DeterministicLevel.INDETERMINATE) { + throw pipelinedUnsupportedError("a statically-indeterminate pipelined producer") + } + // Checksum-mismatch full retry: the runtime counterpart to static indeterminism; it rolls back + // and re-runs succeeding stages on a cross-attempt mismatch, which a group never keeps (moot). + // A PipelinedShuffleDependency does not enable it (see its definition), so this is defensive. + if (shuffleDep.checksumMismatchFullRetryEnabled) { + throw pipelinedUnsupportedError("checksum-mismatch full retry with a pipelined shuffle") + } + // Push-based shuffle merge as the pipelined shuffle: exposes output only after a + // post-completion finalize step, the opposite of incremental reads. A + // PipelinedShuffleDependency disables merge in its constructor, so this is a defensive + // backstop against that being bypassed. + if (shuffleDep.shuffleMergeEnabled) { + throw pipelinedUnsupportedError("push-based shuffle merge as a pipelined shuffle") + } + // A reliable RDD checkpoint in a member's within-stage chain (producer OR consumer side) is + // rejected in checkPipelinedGroupsSupportedInRDDGraph, at job submission before any stage is + // created -- so a reject leaves no partial stage state and both chain sides are covered. + } + /** * We don't support run a barrier stage with dynamic resource allocation enabled, it shall lead * to some confusing behaviors (e.g. with dynamic resource allocation enabled, it may happen that @@ -1060,6 +1181,126 @@ private[spark] class DAGScheduler( false } + /** + * Reject group-level idioms a pipelined group cannot support, checked against the RDD graph + * BEFORE any stage is created -- so a rejection fails the job up front (via handleJobSubmitted's + * listener.jobFailed) without leaving partial scheduler state behind, exactly like the + * speculation check. + * + * Call only for a job that has a pipelined dependency (handleJobSubmitted gates on + * hasPipelined): the resource-profile check below is not keyed on a pipelined dependency, so on a + * regular job it would reject an ordinary RDD.withResources(...) use. Throws + * PIPELINED_SHUFFLE_UNSUPPORTED on violation. Enforces: + * - Fan-out: a pipelined producer feeding more than one consumer. 1:N is a supported model not + * yet built (it needs multicast to N live readers), so it is rejected for now. A + * PipelinedShuffleDependency's producer is `dep.rdd`; a "consumer" is any RDD that lists that + * dependency. More than one distinct consumer RDD for the same pipelined shuffle is fan-out. + * - Reliable RDD checkpoint in a group member's within-stage chain (producer OR consumer side): + * a reliable `checkpoint()` writes a durable, lineage-truncated snapshot, which both + * reintroduces cross-time reuse of a transient edge and requires a post-success recompute + * of the member's transient input -- for a consumer, that input is the vanished pipelined + * shuffle. Checked here (not at stage creation) so a reject leaves no partial stage state, and + * so BOTH the producer chain (rooted at pd.rdd) and each consumer chain (rooted at a consuming + * RDD) are covered from the whole-graph view. + * - A non-default resource profile on any member. The gang slot check compares one demand + * against one profile's capacity and measures it against the default profile, so the whole + * group is required to run on the default profile; any member with an explicit non-default + * profile is rejected. Per-profile accounting is a follow-up. + * + * (The remaining group-level case -- a regular shuffle internal to a group -- is a structural + * invariant that does not arise for the prefix -> pipelined-group -> suffix shapes targeted here: + * groups are split at regular-shuffle boundaries. The producer-side idioms -- barrier, DRA, + * indeterminate, checksum, push-merge -- are rejected in checkPipelinedProducerSupported at stage + * creation, where a producer-only throw leaves no partial state.) + */ + private def checkPipelinedGroupsSupportedInRDDGraph(finalRDD: RDD[_]): Unit = { + // Walk the whole RDD graph once, collecting for each pipelined shuffleId the distinct consumer + // RDDs that read it (for the fan-out check), the producer RDDs that write it (roots of producer + // member stages), and every reliably-checkpointed RDD (to locate ones inside a member stage). + val consumersByShuffleId = new HashMap[Int, HashSet[Int]] + val producerRoots = new HashSet[RDD[_]] // RDDs that WRITE a pipelined shuffle + val reliablyCheckpointed = new HashSet[RDD[_]] // RDDs with a reliable checkpoint pending + var hasNonDefaultResourceProfile = false // any member RDD with a non-default RP + traverseRDDGraph(finalRDD) { (rdd, enqueue) => + if (rdd.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]])) { + reliablyCheckpointed += rdd + } + // An RDD's EFFECTIVE profile is its explicit one, or the default when unset. The slot check + // measures capacity against the default profile (see rejectUnadmittablePipelinedGroup), so + // the whole group must run on the default profile; any explicit non-default profile on a + // member makes the group span profiles (against the default the rest use) and is rejected. + val rp = rdd.getResourceProfile() + if (rp != null && rp.id != ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID) { + hasNonDefaultResourceProfile = true + } + rdd.dependencies.foreach { + case pd: PipelinedShuffleDependency[_, _, _] => + consumersByShuffleId.getOrElseUpdate(pd.shuffleId, new HashSet[Int]) += rdd.id + producerRoots += pd.rdd + enqueue(pd.rdd) + case dep => + enqueue(dep.rdd) + } + } + if (consumersByShuffleId.values.exists(_.size > 1)) { + throw pipelinedUnsupportedError( + "a pipelined producer with more than one consumer (fan-out / branching)") + } + // Resource profile. The gang slot check compares one demand against one profile's capacity + // (maxNumConcurrentTasks is defined per profile) and measures it against the DEFAULT profile, + // so the whole group is required to run on the default profile and fails fast otherwise; + // per-profile accounting is a follow-up. Reject if ANY member carries an explicit non-default + // profile -- that both spans profiles (against the default the other members use) and would be + // admitted against the wrong (default) capacity pool. This is a real check, not just a + // documented assumption: nothing else enforces it. + if (hasNonDefaultResourceProfile) { + throw pipelinedUnsupportedError( + "a pipelined group member with a non-default resource profile (the whole group must run " + + "on the default profile)") + } + // Reject a reliable checkpoint anywhere in a pipelined-group MEMBER's within-stage chain. + // Keyed on checkpointData being ReliableRDDCheckpointData, not isCheckpointed, since the + // write has not happened yet. Cache / .persist() / local checkpoint are whole-partition and + // ephemeral and are not rejected. A reliably-checkpointed RDD `cp` is inside a member stage + // iff: + // - PRODUCER side: cp is within some producer root's own within-stage chain (walk parents from + // the producer root, stopping at shuffle boundaries), OR + // - CONSUMER side: cp's OWN within-stage chain reads a pipelined shuffle (walk parents from + // cp, stopping at shuffle boundaries, and check whether any stopped-at boundary is + // pipelined). + // Rooting the consumer check at each checkpointed RDD (rather than at the PSD-reading RDD) is + // what makes it cover a checkpoint anywhere DOWNSTREAM in the consumer stage, not just on the + // reading RDD itself. + def chainHasReliableCheckpoint(root: RDD[_]): Boolean = + !traverseParentRDDsWithinStage(root, (r: RDD[_]) => + !r.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]])) + val offending = + // CONSUMER side: a checkpointed RDD whose own within-stage chain reads a pipelined shuffle is + // inside a consumer member stage (covers a checkpoint anywhere in that stage, not just on the + // reading RDD). PRODUCER side: a producer root's within-stage chain carries a checkpoint. + reliablyCheckpointed.exists(rddChainReadsPipelinedShuffle) || + producerRoots.exists(chainHasReliableCheckpoint) + if (offending) { + throw pipelinedUnsupportedError( + "a reliable RDD checkpoint in a pipelined-group member's within-stage chain") + } + } + + /** Whether `rdd`'s within-stage chain (parents, stopping at shuffle boundaries) reads through a + * [[PipelinedShuffleDependency]] -- i.e. `rdd` is inside a pipelined CONSUMER member stage. */ + private def rddChainReadsPipelinedShuffle(rdd: RDD[_]): Boolean = { + !traverseRDDGraphUntil(rdd) { (r, enqueue) => + val readsPipelined = r.dependencies.exists { + case _: PipelinedShuffleDependency[_, _, _] => true + case _: ShuffleDependency[_, _, _] => false // regular boundary: not within this stage + case narrowDep => + enqueue(narrowDep.rdd) + false + } + !readsPipelined + } + } + /** Invoke `.partitions` on the given RDD and all of its ancestors */ private def eagerlyComputePartitionsForRddAndAncestors(rdd: RDD[_]): Unit = { val startTime = System.nanoTime @@ -1574,49 +1815,17 @@ private[spark] class DAGScheduler( */ private def rejectUnadmittablePipelinedGroup( jobId: Int, finalRDD: RDD[_], partitions: Array[Int], listener: JobListener): Boolean = { - // Reject up-front (before any stage is created) two group members the admission model cannot - // support, walking the group's RDD graph once: - // - A barrier member. A barrier stage exposes its output only after a global sync, which - // contradicts a pipelined consumer reading the producer's output incrementally as it runs; - // and its failure recovery resubmits the stage, which the group's atomic completion cannot - // accommodate (a resubmitted producer would drop a co-scheduled consumer's buffered - // completions). Reject it here rather than let it be co-scheduled. - // - A member on a non-default resource profile. Admission below measures capacity and - // occupancy against the default profile, but each stage derives its profile from its RDDs - // (see createShuffleMapStage/createResultStage). A member carrying a non-default profile would be - // admitted against the default profile's free slots yet run in a different, often smaller - // pool -- and could then queue or deadlock there. - var offendingBarrier = false - var offendingRp: Option[ResourceProfile] = None - traverseRDDGraph(finalRDD) { (rdd, enqueue) => - if (rdd.isBarrier()) { - offendingBarrier = true - } - val rp = rdd.getResourceProfile() - if (offendingRp.isEmpty && rp != null && rp.id != DEFAULT_RESOURCE_PROFILE_ID) { - offendingRp = Some(rp) - } - rdd.dependencies.foreach(dep => enqueue(dep.rdd)) - } - if (offendingBarrier) { - logWarning(log"Rejecting job ${MDC(JOB_ID, jobId)}: a pipelined stage group contains a " + - log"barrier member") - listener.jobFailed(new SparkException( - "A pipelined shuffle job with a barrier stage in the pipelined group is not supported: a " + - "barrier stage exposes its output only after a global sync, which is incompatible with " + - "a pipelined consumer reading its output incrementally.")) - return true - } - if (offendingRp.nonEmpty) { - logWarning(log"Rejecting job ${MDC(JOB_ID, jobId)}: a pipelined stage group member uses a " + - log"non-default resource profile") - listener.jobFailed(new SparkException( - "A pipelined shuffle job with a member on a non-default resource profile is not " + - s"supported (resource profile id ${offendingRp.get.id}): the whole pipelined stage " + - "group must run on the default resource profile.")) - return true - } if (!sc.conf.get(config.PIPELINED_GROUP_SLOT_CHECK_ENABLED)) { + // The only deadlock-prevention check for gang admission is off. Legitimate only when the + // deployment admits capacity out-of-band (e.g. a slot reservation); otherwise a pipelined + // group that cannot co-fit will be gang-scheduled and can deadlock. Warn once so this is + // never a silent state. + if (!warnedPipelinedSlotCheckDisabled) { + warnedPipelinedSlotCheckDisabled = true + logWarning(log"${MDC(CONFIG, config.PIPELINED_GROUP_SLOT_CHECK_ENABLED.key)}=false: " + + log"pipelined-group gang admission is NOT checking free slots. This is safe only if " + + log"capacity is reserved out-of-band; otherwise a group that cannot co-fit may deadlock.") + } return false } val rp = sc.resourceProfileManager.defaultResourceProfile @@ -1680,24 +1889,12 @@ private[spark] class DAGScheduler( * member's `TaskSet.isPipelined` at submission and routing a member FetchFailed to a whole-group * abort. */ - private def isPipelinedGroupMember(stage: Stage): Boolean = { - if (isPipelinedProducer(stage)) { - return true - } - // Consumer check: does this stage read through a PipelinedShuffleDependency at one of its - // shuffle boundaries? Walk the stage's own RDD graph (descending narrow deps, stopping at every - // shuffle boundary) -- the same edges that define the stage -- and look for a pipelined one. - !traverseRDDGraphUntil(stage.rdd) { (rdd, enqueue) => - val hasPipelinedBoundary = rdd.dependencies.exists { - case _: PipelinedShuffleDependency[_, _, _] => true - case _: ShuffleDependency[_, _, _] => false // regular shuffle boundary: do not descend - case narrowDep => - enqueue(narrowDep.rdd) - false - } - !hasPipelinedBoundary // keep walking until a pipelined boundary is found - } - } + private def isPipelinedGroupMember(stage: Stage): Boolean = + // Producer side: it writes a pipelined shuffle. Consumer side: its within-stage chain reads + // one. rddChainReadsPipelinedShuffle is the single source of truth for that within-stage walk + // (descend narrow deps, stop at every shuffle boundary, look for a pipelined one) -- do not + // re-inline it; the consumer walk here and that method used to be byte-identical copies. + isPipelinedProducer(stage) || rddChainReadsPipelinedShuffle(stage.rdd) /** Finds the earliest-created active job that needs the stage */ // TODO: Probably should actually find among the active jobs that need this @@ -1875,9 +2072,19 @@ private[spark] class DAGScheduler( if (hasPipelined && rejectUnadmittablePipelinedGroup(jobId, finalRDD, partitions, listener)) { return } - var finalStage: ResultStage = null try { + // Reject group-level unsupported pipelined idioms (e.g. fan-out, a non-default resource + // profile, a reliable checkpoint in a member stage) from the RDD graph, up front -- before + // any stage is created, so a rejection leaves no partial scheduler state. Gated on + // hasPipelined: every idiom this checks concerns a pipelined group, so it must not run for a + // job with no pipelined dependency (the resource-profile check in particular is not keyed on + // a pipelined dependency and would otherwise reject an ordinary job that merely uses a + // non-default profile via RDD.withResources). Inside this try so any incidental exception + // from the graph walk is handled by the same listener.jobFailed path as stage creation. + if (hasPipelined) { + checkPipelinedGroupsSupportedInRDDGraph(finalRDD) + } // New stage creation may throw an exception if, for example, jobs are run on a // HadoopRDD whose underlying HDFS files have been deleted. finalStage = createResultStage(finalRDD, func, partitions, jobId, callSite) @@ -1910,6 +2117,15 @@ private[spark] class DAGScheduler( return } + case e: PipelinedShuffleUnsupportedException => + // An up-front idiom rejection (checkPipelinedGroupsSupportedInRDDGraph / a producer-side + // check in createResultStage), not a stage-creation failure. Log it as such (the generic + // "Creating new stage failed" message below would be misleading). Matched by TYPE, not by + // the error-condition string, so a rename or a wrapped cause cannot misroute it. + logWarning(log"Rejecting job ${MDC(JOB_ID, jobId)}: unsupported pipelined-shuffle idiom", e) + listener.jobFailed(e) + return + case e: Exception => logWarning(log"Creating new stage failed due to exception - job: ${MDC(JOB_ID, jobId)}", e) listener.jobFailed(e) @@ -1919,6 +2135,14 @@ private[spark] class DAGScheduler( barrierJobIdToNumTasksCheckFailures.remove(jobId) val job = new ActiveJob(jobId, finalStage, callSite, listener, artifacts, properties) + // Record whether this job uses a pipelined shuffle (computed above), so the per-submit + // pipelined-group checks can short-circuit for a regular job without walking its RDD graph. + job.hasPipelinedDependency = hasPipelined + // We only reach here if the up-front gang admission above (rejectUnadmittablePipelinedGroup) + // did NOT fail-and-return -- i.e. it passed or was disabled by config. Record that so the + // co-schedule path in submitStage can verify the group was admitted, rather than trusting a + // comment. (The job is one group, so admission is job-level; see ActiveJob.) + job.pipelinedGroupAdmitted = hasPipelined clearCacheLocs() logInfo( log"Got job ${MDC(JOB_ID, job.jobId)} (${MDC(CALL_SITE_SHORT_FORM, callSite.shortForm)}) " + @@ -2072,12 +2296,30 @@ private[spark] class DAGScheduler( val allPipelinedParentsRunning = pipelinedMissing.nonEmpty && pipelinedMissing.forall(runningStages.contains) + // INVARIANT (job == one pipelined group): we co-schedule this group's members with + // NO slot check because the whole job was gang-admitted up front in handleJobSubmitted + // (rejectUnadmittablePipelinedGroup) before any member ran. Verify that link rather + // than trusting the comment. This always holds today (the flag is set on the same + // admitted path), so it is an invariant tripwire -- but make it FAIL-CLOSED (abort the + // group, not just assert): a bare `assert` is elided under -da, so if a future change + // relaxes job==group and forgets to move gang admission to a per-group "admit-on-ready" + // check HERE, a prod build would silently gang-schedule an unadmitted group and could + // deadlock. Aborting turns that into a clean, rerunnable job failure. (When that + // per-group check is added, it REPLACES this guard -- do not just delete it.) if (regularMissing.isEmpty && allPipelinedParentsRunning) { + if (!jobIdToActiveJob.get(jobId.get).exists(_.pipelinedGroupAdmitted)) { + abortStage(stage, + s"Co-scheduling pipelined $stage whose job was not gang-admitted up front " + + "(the job==group admission invariant is violated)", None) + return + } // The whole group's capacity was already admitted up front (handleJobSubmitted -> // rejectUnadmittablePipelinedGroup) before any member was submitted, so the group is // known to fit; just co-schedule this consumer with its running producer(s). No slot // check here -- that would re-measure capacity against a mid-flight snapshot and is - // unnecessary once admission is decided up front (gang admission). + // unnecessary once admission is decided up front (gang admission). Group-level + // idiom rejection (fan-out, internal regular shuffle) already happened at job + // submission (checkPipelinedGroupsSupportedInRDDGraph + the all-pipelined check). logInfo(log"Submitting ${MDC(STAGE, stage)} concurrently with its running " + log"pipelined producer(s) ${MDC(MISSING_PARENT_STAGES, pipelinedMissing)}") // Record that this stage is co-scheduled with still-running pipelined producers, @@ -2538,9 +2780,13 @@ private[spark] class DAGScheduler( case _: ResultStage => None } + // Only a job that uses a pipelined shuffle can have a pipelined-group member; gate the + // group-membership graph walk on that cheap per-job flag so a regular job pays nothing here. + val isPipelined = jobIdToActiveJob.get(jobId).exists(_.hasPipelinedDependency) && + isPipelinedGroupMember(stage) taskScheduler.submitTasks(new TaskSet( tasks.toArray, stage.id, stage.latestInfo.attemptNumber(), jobId, properties, - stage.resourceProfileId, shuffleId)) + stage.resourceProfileId, shuffleId, isPipelined = isPipelined)) } else { // Because we posted SparkListenerStageSubmitted earlier, we should mark // the stage as completed here in case there are no tasks to run @@ -3081,7 +3327,32 @@ private[spark] class DAGScheduler( case smt: ShuffleMapTask => val shuffleStage = stage.asInstanceOf[ShuffleMapStage] - if (!ignoreOldTaskAttempts) { + if (shuffleStage.isPipelined) { + // A pipelined shuffle is transient and must NOT be registered with the + // MapOutputTracker as a durable, addressable output: doing so would let + // executor/host loss strip it there, flip isAvailable to false, and resubmit the + // producer -- which then hangs its streaming writer (it blocks forever on + // termination acks from reducers that already finished). Track the completed + // partition locally and monotonically on the stage instead; the incremental shuffle + // reader discovers the producer through its own transport, not the MapOutputTracker. + // Checksum-mismatch detection does not apply (a pipelined dependency never enables + // checksum retry -- see PipelinedShuffleDependency). + // + // Record the partition and decrement pendingPartitions UNCONDITIONALLY -- outside the + // `!ignoreOldTaskAttempts` and bogus-epoch guards that gate a regular shuffle. Both + // guards exist only to avoid trusting a MapOutputTracker registration that a later + // rollback (ignoreOldTaskAttempts, from an indeterminate/rolled-back stage) or an + // executor-loss strip (bogus epoch) could invalidate. A pipelined stage never + // registers there, and its completed set is monotonic and never rolled back (a + // transient shuffle cannot be recomputed; any real group failure aborts the whole + // group). Skipping the record for an "old" or "bogus" straggler would be actively + // harmful: with pendingPartitions decremented but the partition unrecorded, a dropped + // last partition leaves the stage "done but not available" -> processShuffleMapStage- + // Completion resubmits the transient producer, reopening the streaming-writer hang. + // An already-successful straggler is not a failure, so recording it is correct. + shuffleStage.pendingPartitions -= task.partitionId + shuffleStage.addPipelinedCompletedPartition(smt.partitionId) + } else if (!ignoreOldTaskAttempts) { shuffleStage.pendingPartitions -= task.partitionId val status = event.result.asInstanceOf[MapStatus] val execId = status.location.executorId @@ -3149,6 +3420,33 @@ private[spark] class DAGScheduler( log"${MDC(STAGE_ATTEMPT_ID, task.stageAttemptId)} and there is a more recent attempt for " + log"that stage (attempt " + log"${MDC(NUM_ATTEMPT, failedStage.latestInfo.attemptNumber())}) running") + } else if (activeJobForStage(failedStage).flatMap(jobIdToActiveJob.get) + .exists(_.hasPipelinedDependency) && + (isPipelinedGroupMember(failedStage) || isPipelinedGroupMember(mapStage))) { + // Failure is group-atomic for a pipelined group. The base scheduler handles a + // FetchFailed by resubmitting just the map stage in isolation and recomputing serially, + // but a transient pipelined shuffle cannot be re-read and its members are co-scheduled, so + // a lone-stage resubmit is never valid and would deadlock the group. Abort the + // whole group instead: aborting the failed stage tears down its running co-scheduled + // members and fails the job, and the caller (e.g. the streaming batch loop) reruns the + // batch from scratch. This is distinct from the maxTaskFailures=1 lever (which handles + // task failures the TaskSetManager counts): a FetchFailed is NOT counted there (the base + // TaskSetManager marks the task successful and zombies the set), so the routing to group + // failure must be enforced here. + logInfo(log"Failing pipelined group containing ${MDC(FAILED_STAGE, failedStage)} " + + log"(${MDC(FAILED_STAGE_NAME, failedStage.name)}) atomically due to a fetch failure " + + log"from ${MDC(STAGE, mapStage)} (${MDC(STAGE_NAME, mapStage.name)})") + failedStage.failedAttemptIds.add(task.stageAttemptId) + // Still unregister the failed executor's outputs, exactly as the base FetchFailed path + // does -- aborting the group tears down only THIS job's stages, but the FetchFailed is + // authoritative evidence that the executor's shuffle data is gone, and other/concurrent + // jobs sharing that executor must not keep stale MapOutputTracker entries (with an + // external shuffle service, an ExecutorLost would NOT clean these, so this is the only + // proactive channel). Safe for the pipelined shuffle itself: it registers no map outputs + // in the tracker, so this can only strip regular/durable outputs. + unregisterOutputsOnFetchFailedExecutor(bmAddress, task) + abortStage(failedStage, + s"A pipelined group member failed with a fetch failure: $failureMessage", None) } else { val ignoreStageFailure = ignoreDecommissionFetchFailure && isExecutorDecommissioningOrDecommissioned(taskScheduler, bmAddress) @@ -3181,123 +3479,106 @@ private[spark] class DAGScheduler( "longer running") } - if (mapStage.rdd.isBarrier()) { - // Mark all the map as broken in the map stage, to ensure retry all the tasks on - // resubmitted stage attempt. - // TODO: SPARK-35547: Clean all push-based shuffle metadata like merge enabled and - // TODO: finalized as we are clearing all the merge results. - mapOutputTracker.unregisterAllMapAndMergeOutput(shuffleId) - } else if (mapIndex != -1) { - // Mark the map whose fetch failed as broken in the map stage - mapOutputTracker.unregisterMapOutput(shuffleId, mapIndex, bmAddress) - if (pushBasedShuffleEnabled) { - // Possibly unregister the merge result , if the FetchFailed - // mapIndex is part of the merge result of - mapOutputTracker. - unregisterMergeResult(shuffleId, reduceId, bmAddress, Option(mapIndex)) - } + if (mapStage.isPipelined) { + // Defense-in-depth for a pipelined-shuffle FetchFailed that reaches this base path + // rather than the group-atomic branch above (e.g. a job whose hasPipelinedDependency + // flag was not propagated through the group check). The base path is invalid for a + // pipelined shuffle in two ways: (a) the MapOutputTracker invalidation below would + // throw ShuffleStatusNotFoundException (a pipelined shuffle is never registered there, + // see createShuffleMapStage); (b) the resubmit branch would enqueue a lone-stage + // resubmit of the transient producer, which -- as the group-atomic branch's comment + // explains -- is never valid and would deadlock the group. So abort the group here + // instead, matching the group-atomic branch's outcome, and skip both. + abortStage(failedStage, + s"A pipelined group member failed with a fetch failure: $failureMessage", None) } else { - // Unregister the merge result of if there is a FetchFailed event - // and is not a MetaDataFetchException which is signified by bmAddress being null - if (bmAddress != null && - bmAddress.executorId.equals(BlockManagerId.SHUFFLE_MERGER_IDENTIFIER)) { - assert(pushBasedShuffleEnabled, "Push based shuffle expected to " + - "be enabled when handling merge block fetch failure.") - mapOutputTracker. - unregisterMergeResult(shuffleId, reduceId, bmAddress, None) + if (mapStage.rdd.isBarrier()) { + // Mark all the map as broken in the map stage, to ensure retry all the tasks on + // resubmitted stage attempt. + // TODO: SPARK-35547: Clean all push-based shuffle metadata like merge enabled and + // TODO: finalized as we are clearing all the merge results. + mapOutputTracker.unregisterAllMapAndMergeOutput(shuffleId) + } else if (mapIndex != -1) { + // Mark the map whose fetch failed as broken in the map stage + mapOutputTracker.unregisterMapOutput(shuffleId, mapIndex, bmAddress) + if (pushBasedShuffleEnabled) { + // Possibly unregister the merge result , if the FetchFailed + // mapIndex is part of the merge result of + mapOutputTracker. + unregisterMergeResult(shuffleId, reduceId, bmAddress, Option(mapIndex)) + } + } else { + // Unregister the merge result of if there is a FetchFailed + // event and is not a MetaDataFetchException (signified by bmAddress being null) + if (bmAddress != null && + bmAddress.executorId.equals(BlockManagerId.SHUFFLE_MERGER_IDENTIFIER)) { + assert(pushBasedShuffleEnabled, "Push based shuffle expected to " + + "be enabled when handling merge block fetch failure.") + mapOutputTracker. + unregisterMergeResult(shuffleId, reduceId, bmAddress, None) + } } - } - if (failedStage.rdd.isBarrier()) { - failedStage match { - case failedMapStage: ShuffleMapStage => - // Mark all the map as broken in the map stage, to ensure retry all the tasks on - // resubmitted stage attempt. - mapOutputTracker.unregisterAllMapAndMergeOutput(failedMapStage.shuffleDep.shuffleId) - - case failedResultStage: ResultStage => - // Abort the failed result stage since we may have committed output for some - // partitions. - val reason = "Could not recover from a failed barrier ResultStage. Most recent " + - s"failure reason: $failureMessage" - abortStage(failedResultStage, reason, None) + if (failedStage.rdd.isBarrier()) { + failedStage match { + case failedMapStage: ShuffleMapStage => + // Mark all the map as broken in the map stage, to ensure retry all the tasks on + // resubmitted stage attempt. + mapOutputTracker.unregisterAllMapAndMergeOutput( + failedMapStage.shuffleDep.shuffleId) + + case failedResultStage: ResultStage => + // Abort the failed result stage since we may have committed output for some + // partitions. + val reason = "Could not recover from a failed barrier ResultStage. Most recent " + + s"failure reason: $failureMessage" + abortStage(failedResultStage, reason, None) + } } - } - if (shouldAbortStage) { - abortStage(failedStage, abortReason.get, None) - } else { // update failedStages and make sure a ResubmitFailedStages event is enqueued - // TODO: Cancel running tasks in the failed stage -- cf. SPARK-17064 - val noResubmitEnqueued = !failedStages.contains(failedStage) - failedStages += failedStage - failedStages += mapStage - if (noResubmitEnqueued) { - // For statically indeterminate stages, trigger rollback early (here and in - // submitMissingTasks) rather than deferring to task completion. This is more - // efficient because it clears shuffle outputs before the retry is submitted, - // ensuring findMissingPartitions() returns all partitions. - // - // For runtime detection (checksum mismatch), rollback is triggered at task - // completion when the mismatch is discovered. - // - // The `rollbackCurrentStage = true` parameter ensures the failed map stage is - // included in the cleanup: clearing its shuffle outputs, marking old task results - // to be ignored, and creating a new shuffle merge state for the upcoming retry. - if (mapStage.isStaticallyIndeterminate && - !mapStage.shuffleDep.checksumMismatchFullRetryEnabled) { - rollbackSucceedingStages(mapStage, rollbackCurrentStage = true) - } + if (shouldAbortStage) { + abortStage(failedStage, abortReason.get, None) + } else { // update failedStages and make sure a ResubmitFailedStages event is enqueued + // TODO: Cancel running tasks in the failed stage -- cf. SPARK-17064 + val noResubmitEnqueued = !failedStages.contains(failedStage) + failedStages += failedStage + failedStages += mapStage + if (noResubmitEnqueued) { + // For statically indeterminate stages, trigger rollback early (here and in + // submitMissingTasks) rather than deferring to task completion. This is more + // efficient because it clears shuffle outputs before the retry is submitted, + // ensuring findMissingPartitions() returns all partitions. + // + // For runtime detection (checksum mismatch), rollback is triggered at task + // completion when the mismatch is discovered. + // + // The `rollbackCurrentStage = true` parameter ensures the failed map stage is + // included in the cleanup: clearing its shuffle outputs, marking old task results + // to be ignored, and creating a new shuffle merge state for the upcoming retry. + if (mapStage.isStaticallyIndeterminate && + !mapStage.shuffleDep.checksumMismatchFullRetryEnabled) { + rollbackSucceedingStages(mapStage, rollbackCurrentStage = true) + } - // We expect one executor failure to trigger many FetchFailures in rapid succession, - // but all of those task failures can typically be handled by a single resubmission of - // the failed stage. We avoid flooding the scheduler's event queue with resubmit - // messages by checking whether a resubmit is already in the event queue for the - // failed stage. If there is already a resubmit enqueued for a different failed - // stage, that event would also be sufficient to handle the current failed stage, but - // producing a resubmit for each failed stage makes debugging and logging a little - // simpler while not producing an overwhelming number of scheduler events. - logInfo( - log"Resubmitting ${MDC(STAGE, mapStage)} " + - log"(${MDC(STAGE_NAME, mapStage.name)}) and ${MDC(FAILED_STAGE, failedStage)} " + - log"(${MDC(FAILED_STAGE_NAME, failedStage.name)}) due to fetch failure") - scheduleResubmit() + // We expect one executor failure to trigger many FetchFailures in rapid succession, + // but all of those task failures can typically be handled by a single resubmission + // of the failed stage. We avoid flooding the scheduler's event queue with resubmit + // messages by checking whether a resubmit is already in the event queue for the + // failed stage. If there is already a resubmit enqueued for a different failed + // stage, that event would also be sufficient to handle the current failed stage, + // but producing a resubmit for each failed stage makes debugging and logging a + // little simpler while not producing an overwhelming number of scheduler events. + logInfo( + log"Resubmitting ${MDC(STAGE, mapStage)} " + + log"(${MDC(STAGE_NAME, mapStage.name)}) and ${MDC(FAILED_STAGE, failedStage)} " + + log"(${MDC(FAILED_STAGE_NAME, failedStage.name)}) due to fetch failure") + scheduleResubmit() + } } } // TODO: mark the executor as failed only if there were lots of fetch failures on it - if (bmAddress != null) { - val externalShuffleServiceEnabled = env.blockManager.externalShuffleServiceEnabled - val isHostDecommissioned = taskScheduler - .getExecutorDecommissionState(bmAddress.executorId) - .exists(_.workerHost.isDefined) - - // Shuffle output of all executors on host `bmAddress.host` may be lost if: - // - External shuffle service is enabled, so we assume that all shuffle data on node is - // bad. - // - Host is decommissioned, thus all executors on that host will die. - val shuffleOutputOfEntireHostLost = externalShuffleServiceEnabled || - isHostDecommissioned - val hostToUnregisterOutputs = if (shuffleOutputOfEntireHostLost - && unRegisterOutputOnHostOnFetchFailure) { - Some(bmAddress.host) - } else { - // Unregister shuffle data just for one executor (we don't have any - // reason to believe shuffle data has been lost for the entire host). - None - } - removeExecutorAndUnregisterOutputs( - execId = bmAddress.executorId, - fileLost = true, - hostToUnregisterOutputs = hostToUnregisterOutputs, - maybeEpoch = Some(task.epoch), - // shuffleFileLostEpoch is ignored when a host is decommissioned because some - // decommissioned executors on that host might have been removed before this fetch - // failure and might have bumped up the shuffleFileLostEpoch. We ignore that, and - // proceed with unconditional removal of shuffle outputs from all executors on that - // host, including from those that we still haven't confirmed as lost due to heartbeat - // delays. - ignoreShuffleFileLostEpoch = isHostDecommissioned) - } + unregisterOutputsOnFetchFailedExecutor(bmAddress, task) } case failure: TaskFailedReason if task.isBarrier => @@ -3837,6 +4118,52 @@ private[spark] class DAGScheduler( maybeEpoch = None) } + /** + * On a FetchFailed, unregister the shuffle outputs of the executor (or its whole host) whose + * fetch failed, treating the FetchFailed as authoritative evidence that its shuffle data is gone. + * Extracted from the base FetchFailed handler so the pipelined-group-abort branch can also run it: + * aborting the group fails only this job's stages, but a dead executor's REGULAR outputs must + * still be cleaned up for other/concurrent jobs (with an external shuffle service, an ExecutorLost + * does not clean them, so FetchFailed is the only proactive channel). No-op when `bmAddress` is + * null. Safe for a pipelined shuffle: it registers no map outputs in the tracker, so this can only + * strip regular/durable outputs. + */ + private def unregisterOutputsOnFetchFailedExecutor( + bmAddress: BlockManagerId, task: Task[_]): Unit = { + // TODO: mark the executor as failed only if there were lots of fetch failures on it + if (bmAddress != null) { + val externalShuffleServiceEnabled = env.blockManager.externalShuffleServiceEnabled + val isHostDecommissioned = taskScheduler + .getExecutorDecommissionState(bmAddress.executorId) + .exists(_.workerHost.isDefined) + + // Shuffle output of all executors on host `bmAddress.host` may be lost if: + // - External shuffle service is enabled, so we assume that all shuffle data on node is bad. + // - Host is decommissioned, thus all executors on that host will die. + val shuffleOutputOfEntireHostLost = externalShuffleServiceEnabled || isHostDecommissioned + val hostToUnregisterOutputs = if (shuffleOutputOfEntireHostLost + && unRegisterOutputOnHostOnFetchFailure) { + Some(bmAddress.host) + } else { + // Unregister shuffle data just for one executor (we don't have any + // reason to believe shuffle data has been lost for the entire host). + None + } + removeExecutorAndUnregisterOutputs( + execId = bmAddress.executorId, + fileLost = true, + hostToUnregisterOutputs = hostToUnregisterOutputs, + maybeEpoch = Some(task.epoch), + // shuffleFileLostEpoch is ignored when a host is decommissioned because some + // decommissioned executors on that host might have been removed before this fetch + // failure and might have bumped up the shuffleFileLostEpoch. We ignore that, and + // proceed with unconditional removal of shuffle outputs from all executors on that + // host, including from those that we still haven't confirmed as lost due to heartbeat + // delays. + ignoreShuffleFileLostEpoch = isHostDecommissioned) + } + } + /** * Handles removing an executor from the BlockManagerMaster as well as unregistering shuffle * outputs for the executor or optionally its host. @@ -4449,6 +4776,20 @@ private[spark] object DAGScheduler { val RESUBMIT_TIMEOUT = 200 } +/** + * Thrown when a job uses a pipelined-shuffle idiom that is not supported (fan-out, a barrier / + * indeterminate / checksum-retry / push-merge producer, a reliable checkpoint in a member's chain, + * or a non-default resource profile on a member). `handleJobSubmitted` matches on this + * TYPE to distinguish an up-front idiom rejection from an ordinary stage-creation failure, not on + * the error-condition string (which a rename or a wrapped cause would silently break). Carries the + * `PIPELINED_SHUFFLE_UNSUPPORTED` error class so the user-facing message is unchanged. + */ +private[scheduler] class PipelinedShuffleUnsupportedException(reason: String) + extends SparkException( + errorClass = "PIPELINED_SHUFFLE_UNSUPPORTED", + messageParameters = scala.collection.immutable.Map("reason" -> reason), + cause = null) + /** * A NOT thread-safe set that only keeps the last `capacity` elements added to it. */ diff --git a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala index 79f7af48f102a..135e5e5a4bf2c 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala @@ -19,7 +19,7 @@ package org.apache.spark.scheduler import scala.collection.mutable.HashSet -import org.apache.spark.{MapOutputTrackerMaster, ShuffleDependency} +import org.apache.spark.{MapOutputTrackerMaster, PipelinedShuffleDependency, ShuffleDependency} import org.apache.spark.rdd.{DeterministicLevel, RDD} import org.apache.spark.util.CallSite @@ -59,6 +59,33 @@ private[spark] class ShuffleMapStage( */ val pendingPartitions = new HashSet[Int] + /** Whether this stage produces a pipelined (incrementally-readable) shuffle. */ + val isPipelined: Boolean = shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]] + + /** + * Availability tracking for a pipelined shuffle. A pipelined shuffle produces no durable, + * addressable map output and is NOT registered with the `MapOutputTracker` at all -- + * `createShuffleMapStage` registers it only with the `StreamingShuffleOutputTracker` (the two + * trackers are split by dependency type, with no overlap). Its map-stage availability therefore + * cannot be read from the `MapOutputTracker`; instead we track completed partitions here, + * monotonically: a partition is added when its map task succeeds and is NEVER removed on + * executor/host loss. + * + * This is the crux of avoiding the streaming-writer resubmit hang: if a pipelined shuffle's + * availability were read from the `MapOutputTracker`, losing an executor that held a completed + * (already-consumed) pipelined output would strip it there, flip `isAvailable` to false, and make + * the DAGScheduler resubmit the producer -- whose streaming writer then blocks forever waiting + * for termination acks from reducers that already finished. Keeping availability local and + * monotonic means executor loss never triggers such a resubmit; a genuine mid-group failure is + * handled group-atomically instead. Unused (and empty) for a non-pipelined stage. + */ + private[this] val pipelinedCompletedPartitions = new HashSet[Int] + + /** Record a successful map task's partition as completed (pipelined stages only). */ + private[scheduler] def addPipelinedCompletedPartition(partitionId: Int): Unit = { + pipelinedCompletedPartitions += partitionId + } + override def toString: String = "ShuffleMapStage " + id /** @@ -81,7 +108,12 @@ private[spark] class ShuffleMapStage( * Number of partitions that have shuffle outputs. * When this reaches [[numPartitions]], this map stage is ready. */ - def numAvailableOutputs: Int = mapOutputTrackerMaster.getNumAvailableOutputs(shuffleDep.shuffleId) + def numAvailableOutputs: Int = { + // A pipelined shuffle's outputs are never populated in the MapOutputTracker (its entry stays + // empty; see pipelinedCompletedPartitions); read its locally-tracked, monotonic completed set. + if (isPipelined) pipelinedCompletedPartitions.size + else mapOutputTrackerMaster.getNumAvailableOutputs(shuffleDep.shuffleId) + } /** * Returns true if the map stage is ready, i.e. all partitions have shuffle outputs. @@ -90,9 +122,13 @@ private[spark] class ShuffleMapStage( /** Returns the sequence of partition ids that are missing (i.e. needs to be computed). */ override def findMissingPartitions(): Seq[Int] = { - mapOutputTrackerMaster - .findMissingPartitions(shuffleDep.shuffleId) - .getOrElse(0 until numPartitions) + if (isPipelined) { + (0 until numPartitions).filterNot(pipelinedCompletedPartitions.contains) + } else { + mapOutputTrackerMaster + .findMissingPartitions(shuffleDep.shuffleId) + .getOrElse(0 until numPartitions) + } } /** diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSet.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSet.scala index 3513cb1f93764..407ebc53fcffc 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskSet.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSet.scala @@ -33,7 +33,13 @@ private[spark] class TaskSet( val priority: Int, val properties: Properties, val resourceProfileId: Int, - val shuffleId: Option[Int]) { + val shuffleId: Option[Int], + // True if this stage is a member of a pipelined group (connected to another stage by a + // PipelinedShuffleDependency). Such a stage's transient shuffle output cannot be re-read in + // isolation, so any task failure must fail the whole group rather than be retried per-task; + // the TaskSetManager uses this to fail fast (maxTaskFailures = 1) and to count every failure, + // including otherwise-uncounted ones like executor loss. Defaults to false. + val isPipelined: Boolean = false) { val id: String = s"$stageId.$stageAttemptId" override def toString: String = "TaskSet " + id diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala index 3a5f60ae19358..761a0d350529c 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala @@ -70,6 +70,15 @@ private[spark] class TaskSetManager( val ser = env.closureSerializer.newInstance() val tasks = taskSet.tasks + + // A pipelined-group member reads/writes a transient shuffle that cannot be re-read in isolation, + // so a single task failure must fail the whole group (which fails the job, triggering a rerun) + // rather than be retried per-task. For such a task set we cap attempts at 1 and count every + // failure, including reasons that are normally not counted (e.g. executor loss). This is the + // native equivalent of the prototype's per-batch "any failure counts" behavior, keyed on the + // pipelined dependency (via TaskSet.isPipelined) rather than a job property. + private val effectiveMaxTaskFailures = if (taskSet.isPipelined) 1 else maxTaskFailures + private val isShuffleMapTasks = tasks(0).isInstanceOf[ShuffleMapTask] // shuffleId is only available when isShuffleMapTasks=true private val shuffleId = taskSet.shuffleId @@ -1106,16 +1115,31 @@ private[spark] class TaskSetManager( emptyTaskInfoAccumulablesAndNotifyDagScheduler(tid, tasks(index), reason, null, accumUpdates, metricPeaks) - if (!isZombie && reason.countTowardsTaskFailures) { + // A pipelined-group member's transient shuffle cannot be recovered per-task, so a GENUINE task + // failure must fail the whole group even when the reason is normally not counted -- most + // importantly executor loss not caused by the app (ExecutorLostFailure with + // exitCausedByApp=false), which strands the transient output. But we must NOT force-count + // reasons that are benign by design: TaskKilled (a losing speculative/duplicate attempt after + // another attempt succeeded, or a deliberate kill) and TaskCommitDenied (a speculation commit + // race) are not real failures, and counting them would abort the group spuriously. So for a + // pipelined set we count everything EXCEPT those two benign reasons. Non-pipelined sets are + // unchanged. + val forceCountForPipelined = taskSet.isPipelined && (reason match { + case _: TaskKilled | _: TaskCommitDenied => false + case _ => true + }) + val countTowardsTaskFailures = reason.countTowardsTaskFailures || forceCountForPipelined + if (!isZombie && countTowardsTaskFailures) { assert (null != failureReason) taskSetExcludelistHelperOpt.foreach(_.updateExcludedForFailedTask( info.host, info.executorId, index, failureReasonString)) numFailures(index) += 1 - if (numFailures(index) >= maxTaskFailures) { + if (numFailures(index) >= effectiveMaxTaskFailures) { logError(log"Task ${MDC(TASK_INDEX, index)} in stage " + taskSet.logId + - log" failed ${MDC(MAX_ATTEMPTS, maxTaskFailures)} times; aborting job") + log" failed ${MDC(MAX_ATTEMPTS, effectiveMaxTaskFailures)} times; aborting job") abort("Task %d in stage %s failed %d times, most recent failure: %s\nDriver stacktrace:" - .format(index, taskSet.id, maxTaskFailures, failureReasonString), failureException) + .format(index, taskSet.id, effectiveMaxTaskFailures, failureReasonString), + failureException) return } } @@ -1191,7 +1215,16 @@ 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 && + // A pipelined-group member must never single-resubmit an already-successful map task: its + // transient shuffle output is not addressable and a lone producer rerun hangs the streaming + // writer in awaitTerminationAcks. This "Resubmitted" re-enqueue loop bypasses handleFailedTask, + // so the group-atomic abort (driven from handleFailedTask via maxTaskFailures=1) would never + // see it; exclude pipelined sets here. A genuine executor loss still flows through the + // running-task loop below (iter2 -> handleFailedTask(ExecutorLostFailure)), force-counted for a + // 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. + val maybeShuffleMapOutputLoss = isShuffleMapTasks && !taskSet.isPipelined && !sched.sc.shuffleDriverComponents.supportsReliableStorage() && (reason.isInstanceOf[ExecutorDecommission] || !env.blockManager.externalShuffleServiceEnabled) if (maybeShuffleMapOutputLoss && !isZombie) { diff --git a/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManager.scala b/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManager.scala index 03ca257353c59..b83bd33552b90 100644 --- a/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManager.scala +++ b/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManager.scala @@ -113,8 +113,8 @@ private[spark] class StreamingShuffleManager extends PipelinedShuffleManager wit override def unregisterShuffle(shuffleId: Int): Boolean = { // No manager-side state to release here: the driver's StreamingShuffleOutputTracker is - // unregistered in BlockManagerStorageEndpoint's RemoveShuffle handler, and per-task writer - // and reader resources are released via task completion listeners. + // unregistered in ContextCleaner.doCleanupShuffle (its state is driver-only), and per-task + // writer and reader resources are released via task completion listeners. true } diff --git a/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala b/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala index bb2d7d5c4d8e4..e8f33b7803da9 100644 --- a/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala +++ b/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala @@ -20,6 +20,7 @@ package org.apache.spark import scala.collection.mutable.HashSet import scala.util.Random +import org.apache.logging.log4j.Level import org.scalatest.BeforeAndAfter import org.scalatest.concurrent.Eventually._ import org.scalatest.concurrent.PatienceConfiguration @@ -127,6 +128,52 @@ class ContextCleanerSuite extends ContextCleanerSuiteBase { assert(rdd.collect().toList.equals(collected)) } + test("cleanup shuffle unregisters a pipelined shuffle from the streaming output tracker") { + // A pipelined shuffle lives ONLY in the driver-only StreamingShuffleOutputTracker -- it is + // never registered with the MapOutputTracker (see DAGScheduler.createShuffleMapStage). So the + // MapOutputTracker.containsShuffle guard is false for it, and doCleanupShuffle must still clean + // it up via its own streaming branch; otherwise the tracker grows without bound across + // micro-batches in Real-Time Mode. + val streamingTracker = sc.env.streamingShuffleOutputTracker.get + .asInstanceOf[StreamingShuffleOutputTrackerMaster] + val mapOutputTracker = sc.env.mapOutputTracker.asInstanceOf[MapOutputTrackerMaster] + + // Register a shuffle id in the streaming tracker ONLY, exactly as the scheduler now does for a + // pipelined shuffle (nothing is put in the MapOutputTracker). + val shuffleId = 1000 + streamingTracker.registerShuffle(shuffleId, numMaps = 2, numReduces = 2, jobId = 0) + assert(streamingTracker.containsShuffle(shuffleId)) + assert(!mapOutputTracker.containsShuffle(shuffleId), + "a pipelined shuffle must not be registered with the MapOutputTracker") + + cleaner.doCleanupShuffle(shuffleId, blocking = true) + + // The streaming tracker entry is gone -- cleanup fired even though the shuffle was never in the + // MapOutputTracker, proving the streaming cleanup branch is independent of MapOutputTracker. + assert(!streamingTracker.containsShuffle(shuffleId)) + } + + test("cleanup shuffle leaves a regular shuffle's streaming tracker untouched and quiet") { + // A regular (non-pipelined) shuffle is registered only with the MapOutputTracker, never the + // streaming tracker. doCleanupShuffle must take the MapOutputTracker branch and never touch the + // streaming tracker -- in particular it must not emit the streaming tracker's "attempting to + // unregister a shuffle that hasn't been registered" warning. + val streamingTracker = sc.env.streamingShuffleOutputTracker.get + .asInstanceOf[StreamingShuffleOutputTrackerMaster] + + val (rdd, shuffleDeps) = newRDDWithShuffleDependencies() + rdd.collect() + shuffleDeps.foreach(s => assert(!streamingTracker.containsShuffle(s.shuffleId))) + + val logAppender = new LogAppender("unregister streaming shuffle") + logAppender.setThreshold(Level.WARN) + withLogAppender(logAppender, level = Some(Level.WARN)) { + shuffleDeps.foreach(s => cleaner.doCleanupShuffle(s.shuffleId, blocking = true)) + } + assert(!logAppender.loggingEvents.exists( + _.getMessage.getFormattedMessage.contains("hasn't been registered"))) + } + test("cleanup broadcast") { val broadcast = newBroadcast() val tester = new CleanerTester(sc, broadcastIds = Seq(broadcast.id)) diff --git a/core/src/test/scala/org/apache/spark/StreamingShuffleOutputTrackerSuite.scala b/core/src/test/scala/org/apache/spark/StreamingShuffleOutputTrackerSuite.scala index 3c8d9b2f14e21..a9351630ad8db 100644 --- a/core/src/test/scala/org/apache/spark/StreamingShuffleOutputTrackerSuite.scala +++ b/core/src/test/scala/org/apache/spark/StreamingShuffleOutputTrackerSuite.scala @@ -270,6 +270,20 @@ class StreamingShuffleOutputTrackerSuite assert(tracker.getAllShuffleWriterTaskLocations(0).isEmpty) } + test("StreamingShuffleOutputTrackerMaster - containsShuffle reflects register/unregister") { + val conf = new SparkConf(false) + val tracker = newTrackerMaster(conf) + + // A shuffle that was never registered is absent. + assert(!tracker.containsShuffle(0)) + + tracker.registerShuffle(shuffleId = 0, numMaps = 1, numReduces = 1, jobId = 1) + assert(tracker.containsShuffle(0)) + + tracker.unregisterShuffle(0) + assert(!tracker.containsShuffle(0)) + } + test("StreamingShuffleOutputTrackerMaster - register writer before shuffle fails") { val conf = new SparkConf(false) val tracker = newTrackerMaster(conf) diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index f43354c808652..2c3ad57f8fa5c 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -41,7 +41,7 @@ import org.apache.spark.executor.ExecutorMetrics import org.apache.spark.internal.config import org.apache.spark.internal.config.{LEGACY_ABORT_STAGE_AFTER_KILL_TASKS, Tests} import org.apache.spark.network.shuffle.ExternalBlockStoreClient -import org.apache.spark.rdd.{DeterministicLevel, RDD} +import org.apache.spark.rdd.{DeterministicLevel, RDD, ReliableRDDCheckpointData} import org.apache.spark.resource.{ExecutorResourceRequests, ResourceProfile, ResourceProfileBuilder, TaskResourceProfile, TaskResourceRequests} import org.apache.spark.resource.ResourceUtils.{FPGA, GPU} import org.apache.spark.rpc.RpcTimeoutException @@ -6432,6 +6432,31 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } + test("pipelined shuffle: a regular-shuffle prefix feeding a pipelined producer is rejected") { + // Also mixed: a regular shuffle in the PREFIX that feeds a pipelined producer + // (regularRoot --regular--> producer(pipelined) --pipelined--> consumer). Rather than treat + // the regular edge as an ordinary external input to the group and support a mid-DAG regular + // prefix, a pipelined job rejects ANY regular shuffle up front (the supported shape is + // scan-of-files --pipelined--> stateful, with no upstream shuffle). + val regularRoot = new MyRDD(sc, 2, Nil) + val regularDep = new ShuffleDependency(regularRoot, new HashPartitioner(2)) + val producerRdd = new MyRDD(sc, 2, List(regularDep), tracker = mapOutputTracker) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = {} + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(consumerRdd, Array(0, 1), listener = failListener) + + assert(failure.get() != null, "a pipelined job with a regular-shuffle prefix must fail") + assert(failure.get().getMessage.contains("all-regular or all-pipelined"), + s"expected a mixed-job rejection, got: ${failure.get().getMessage}") + assert(taskSets.isEmpty, "no stage should be submitted for a rejected mixed job") + assertDataStructuresEmpty() + } + test("pipelined shuffle: deep chain A->B->C is submitted fully concurrently") { // A --pipelined--> B --pipelined--> C : all three co-scheduled (each edge is non-sequencing). val rddA = new MyRDD(sc, 2, Nil) @@ -6459,6 +6484,35 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } + test("pipelined shuffle: an explicit job cancellation cleans up a buffered consumer deferral") { + // A buffered deferral is the only mutable state this feature adds; it must never outlive its + // job. Job cancellation goes through failJobAndIndependentStages -> + // cleanupStateForJobAndIndependentStages, which drops the entry (as a consumer key) and removes + // the stage from every other consumer's pending-producer set. Verify a consumer whose + // completion is buffered while its producer runs leaves NO deferral behind when the job is + // cancelled, and no buffered success is later applied as a result. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + val jobId = submit(consumerRdd, Array(0, 1)) + val consumerTaskSet = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + + // Consumer finishes ahead of its producer -> its completion is buffered (a deferral exists). + complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) + assert(scheduler.dependentStageMap.nonEmpty, "a deferral must exist while the producer runs") + assert(results.isEmpty, "the consumer result must be deferred, not applied yet") + + // Cancel the job. The deferral must be torn down with everything else -- no stale entry, and + // the buffered success must not be replayed as a result. + cancel(jobId) + assert(scheduler.dependentStageMap.isEmpty, + "job cancellation must clean up the buffered consumer deferral (no state outlives the job)") + assert(results.isEmpty, "a cancelled job's buffered consumer success must not be applied") + assertDataStructuresEmpty() + } + test("regular shuffle job with speculation enabled is NOT rejected (rejection path is inert)") { // The speculation fail-fast must apply only to jobs with a pipelined dependency; a plain // regular-shuffle job with speculation on runs normally. @@ -6532,65 +6586,6 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } } - test("pipelined shuffle: a member on a non-default resource profile is rejected up front") { - // Admission measures capacity/occupancy against the DEFAULT resource profile, but a stage - // derives its profile from its RDDs. A pipelined group member on a custom profile would be - // admitted against the default pool's free slots yet run in a different (often smaller) pool - // and could deadlock there. Reject such a job before any stage is created. - // Ensure the default profile exists (id 0) before building a custom one, so the custom profile - // gets a distinct, non-default id regardless of test-execution order (profile ids come from a - // process-wide counter, and the default profile occupies id 0). - sc.resourceProfileManager.defaultResourceProfile - val ereqs = new ExecutorResourceRequests().cores(4) - val treqs = new TaskResourceRequests().cpus(2) - val customRp = new ResourceProfileBuilder().require(ereqs).require(treqs).build() - val producerRdd = new MyRDD(sc, 2, Nil) - val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) - val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) - .withResources(customRp) - // Sanity: the consumer really carries a non-default profile (otherwise the test is vacuous). - assert(consumerRdd.getResourceProfile() != null && - consumerRdd.getResourceProfile().id != ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID, - s"test setup: expected a non-default profile, got id " + - s"${Option(consumerRdd.getResourceProfile()).map(_.id)}") - val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() - val failListener = new JobListener { - override def taskSucceeded(index: Int, result: Any): Unit = {} - override def jobFailed(exception: Exception): Unit = failure.set(exception) - } - submit(consumerRdd, Array(0, 1), listener = failListener) - assert(failure.get() != null, - "a pipelined job with a non-default-resource-profile member should be rejected") - assert(failure.get().getMessage.contains("non-default resource profile")) - assert(taskSets.isEmpty, "no stage should be created for a rejected job") - assertDataStructuresEmpty() - } - - test("pipelined shuffle: a barrier group member is rejected up front") { - // A barrier stage exposes its output only after a global sync, which is incompatible with a - // pipelined consumer reading the producer's output incrementally. It is also a recovery hazard: - // a barrier task failure fails the stage and RESUBMITS it (markStageAsFinished without - // willRetry), and if such a producer were co-scheduled, the resubmit would drop a fan-in - // consumer's buffered completions and the job could never complete. Reject any job whose - // pipelined group contains a barrier member before any stage is created, so that path is - // unreachable. - val producerRdd = new MyRDD(sc, 2, Nil).barrier().mapPartitions(iter => iter) - assert(producerRdd.isBarrier(), "test setup: the producer must be a barrier RDD") - val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) - val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) - val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() - val failListener = new JobListener { - override def taskSucceeded(index: Int, result: Any): Unit = {} - override def jobFailed(exception: Exception): Unit = failure.set(exception) - } - submit(consumerRdd, Array(0, 1), listener = failListener) - assert(failure.get() != null, - "a pipelined job with a barrier group member should be rejected") - assert(failure.get().getMessage.contains("barrier")) - assert(taskSets.isEmpty, "no stage should be created for a rejected job") - assertDataStructuresEmpty() - } - test("regular shuffle job with dynamic allocation enabled is NOT rejected (path is inert)") { // The dynamic-allocation fail-fast must apply only to jobs with a pipelined dependency; a plain // regular-shuffle job with dynamic allocation on runs normally. @@ -6832,13 +6827,19 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } } - test("pipelined shuffle: a diamond reusing one producer counts its tasks once for admission") { + test("pipelined shuffle: admission counts a reused producer once (diamond rejected for " + + "fan-out, not for slots)") { // Fan-out/diamond: ONE PipelinedShuffleDependency (one producer, one shuffle id) is read by TWO - // consumers that are then narrow-joined into the result. Execution creates ONE producer stage - // (getOrCreateShuffleMapStage keys on shuffle id), so the group's real concurrent demand is - // producer(2) + result(2) = 4 -- NOT 6. Counting the producer once per consumer EDGE would make - // demand 6 and wrongly reject this group at capacity 4. Pinning capacity to exactly 4 admits - // the correctly-deduped group and would fail if the producer were double-counted. + // consumers that are then narrow-joined into the result. Fan-out is unsupported, so the group + // is ultimately rejected -- but WHICH rejection it gets proves the admission demand is deduped. + // The up-front slot admission (rejectUnadmittablePipelinedGroup) runs BEFORE the fan-out idiom + // check (checkPipelinedGroupsSupportedInRDDGraph). Execution creates ONE producer stage + // (getOrCreateShuffleMapStage keys on shuffle id), so the real concurrent demand is + // producer(2) + result(2) = 4. Counting the producer once per consumer EDGE would inflate it + // to 6. Pinning capacity to exactly 4: the deduped group PASSES the slot check and is then + // rejected for fan-out (PIPELINED_SHUFFLE_UNSUPPORTED); a double-counted group (6 > 4) would + // instead be rejected earlier for CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT. So a dedup regression + // flips the error class -- which this test detects. val producerRdd = new MyRDD(sc, 2, Nil) val myScheduler = scheduler.asInstanceOf[MyDAGScheduler] myScheduler.maxConcurrentTasksForTest = 4 @@ -6852,13 +6853,13 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti val resultRdd = new MyRDD( sc, 2, List(new OneToOneDependency(consumer1), new OneToOneDependency(consumer2)), tracker = mapOutputTracker) - submit(resultRdd, Array(0, 1)) - assert(taskSets.size === 2, - s"diamond group is producer + result (both consumers); got ${taskSets.size} task sets") - // Drain to completion: producer stage first, then the result stage. - completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2) - complete(taskSets(1), Seq((Success, 42), (Success, 43))) - assert(results === Map(0 -> 42, 1 -> 43)) + val failure = submitAndCaptureFailure(resultRdd, Array(0, 1)) + // Rejected for fan-out (demand fit within capacity 4), NOT for slots -- proving the reused + // producer was counted once. + assertPipelinedUnsupported(failure, "more than one consumer") + assert(!failure.getMessage.contains("CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT"), + s"a reused producer must be counted once (fit capacity 4), got a slot rejection: " + + s"${failure.getMessage}") assertDataStructuresEmpty() } finally { myScheduler.maxConcurrentTasksForTest = 1000 @@ -6866,10 +6867,13 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } } - test("pipelined shuffle: a wide fan-out reusing one producer counts its tasks once") { + test("pipelined shuffle: admission counts a reused producer once across a wide fan-out " + + "(rejected for fan-out, not for slots)") { // Wider fan-out: ONE producer feeds THREE consumers, all narrow-joined into the result. Real - // demand is still producer(2) + result(2) = 4; a per-edge count would be 2 + 3*2 = 8. Capacity - // 4 admits the deduped group and would reject the double-counted one. + // demand is still producer(2) + result(2) = 4; a per-edge count would be 2 + 3*2 = 8. At + // capacity 4 the deduped group passes the slot check and is rejected for fan-out; a + // double-counted group (8 > 4) would be rejected first for insufficient slots. See the diamond + // test above for the ordering rationale. val producerRdd = new MyRDD(sc, 2, Nil) val myScheduler = scheduler.asInstanceOf[MyDAGScheduler] myScheduler.maxConcurrentTasksForTest = 4 @@ -6880,12 +6884,11 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti new OneToOneDependency(new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)) }.toList val resultRdd = new MyRDD(sc, 2, consumers, tracker = mapOutputTracker) - submit(resultRdd, Array(0, 1)) - assert(taskSets.size === 2, - s"wide fan-out group is producer + result; got ${taskSets.size} task sets") - completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2) - complete(taskSets(1), Seq((Success, 42), (Success, 43))) - assert(results === Map(0 -> 42, 1 -> 43)) + val failure = submitAndCaptureFailure(resultRdd, Array(0, 1)) + assertPipelinedUnsupported(failure, "more than one consumer") + assert(!failure.getMessage.contains("CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT"), + s"a reused producer must be counted once (fit capacity 4), got a slot rejection: " + + s"${failure.getMessage}") assertDataStructuresEmpty() } finally { myScheduler.maxConcurrentTasksForTest = 1000 @@ -7075,6 +7078,45 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti s"a dropped fan-in consumer must not replay when a surviving producer succeeds; got $results") } + test("pipelined shuffle: a multi-producer consumer's buffered successes are dropped when one " + + "producer fails (no stale replay via a sibling)") { + // Group-atomic drop across MULTIPLE producers. Consumer C reads TWO pipelined producers P1 and + // P2 (a join; purely all-pipelined, no regular shuffle). C finishes early -> its successes are + // buffered against BOTH producers. If P1 then fails, the whole group must be torn down and C's + // successes DROPPED -- they depended on P1's (now-invalid) output. Note the release path + // evaluates producerFailed per finishing producer, so a naive impl could remove P1 + // (parents still has P2, no drop), then later see P2 "succeed" and REPLAY. This asserts the + // shipped stack drops instead: a member failure aborts the whole group, which tears C down. + val rddP1 = new MyRDD(sc, 2, Nil) + val psdP1 = new PipelinedShuffleDependency(rddP1, new HashPartitioner(2)) + val rddP2 = new MyRDD(sc, 2, Nil) + val psdP2 = new PipelinedShuffleDependency(rddP2, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(psdP1, psdP2), tracker = mapOutputTracker) + val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = results.put(index, result) + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(consumerRdd, Array(0, 1), listener = failListener) + // P1, P2, and C are all co-scheduled (purely all-pipelined join, admitted up front). + assert(taskSets.size === 3, s"expected P1, P2, consumer co-scheduled, got ${taskSets.size}") + val tsP1 = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddP1).get + val tsC = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd).get + + // Consumer finishes early -> buffered against both P1 and P2. + complete(tsC, Seq((Success, 42), (Success, 43))) + assert(results.isEmpty, "consumer completions should be buffered while producers run") + assert(scheduler.dependentStageMap.keys.exists(_.rdd eq consumerRdd), + "consumer should be deferred against its two producers") + + // P1 fails. The group is torn down; C's buffered successes must be DROPPED, never replayed -- + // even though P2 has not (and now will not) complete. + failed(tsP1, "producer P1 blew up") + assert(failure.get() != null, "the job must fail when a pipelined producer fails") + assert(results.isEmpty, + "a multi-producer consumer's buffered successes must be dropped when any producer fails") + assertDataStructuresEmpty() + } test("pipelined shuffle: a deferred consumer task fires its TaskEnd exactly once (at replay)") { // A deferred CompletionEvent must have its side effects (task-end listener event, accumulator @@ -7219,6 +7261,825 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } } + test("pipelined shuffle: both producer and consumer task sets are marked isPipelined") { + // The TaskSet.isPipelined flag drives group-atomic failure in the task scheduler; verify the + // DAGScheduler sets it for both members of a pipelined group (producer and consumer). + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + assert(taskSets.size === 2) + val producerTs = + taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq producerRdd).get + val consumerTs = + taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd).get + assert(producerTs.isPipelined, "the pipelined producer's task set must be marked isPipelined") + assert(consumerTs.isPipelined, "the pipelined consumer's task set must be marked isPipelined") + + completeShuffleMapStageSuccessfully(producerTs.stageId, 0, 2) + complete(consumerTs, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("regular shuffle: task sets are NOT marked isPipelined (inertness)") { + // A regular producer/consumer must not be marked isPipelined. + val producerRdd = new MyRDD(sc, 2, Nil) + val regularDep = new ShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(regularDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + assert(taskSets.head.isPipelined === false, "a regular producer must not be marked isPipelined") + completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2) + assert(taskSets(1).isPipelined === false, "a regular consumer must not be marked isPipelined") + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + // ========================================================================================== + // Cross-job / cross-time reuse prevention at both layers: a consumed pipelined + // producer whose executor is lost must not be resubmitted + // ========================================================================================== + + test("pipelined shuffle: producer availability is tracked on the stage, not the" + + "MapOutputTracker") { + // A pipelined producer's completed partitions are tracked on ShuffleMapStage (monotonic), not + // registered as durable outputs in the MapOutputTracker. Verify: after the producer's map tasks + // succeed, the stage is available WITHOUT the shuffle having map outputs in the tracker. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val shuffleId = pipelinedDep.shuffleId + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + val producerStageId = taskSets.head.stageId + val producerStage = + scheduler.stageIdToStage(producerStageId).asInstanceOf[ShuffleMapStage] + + // Complete the producer's two map tasks. + complete(taskSets.head, Seq( + (Success, makeMapStatus("hostA", 2)), + (Success, makeMapStatus("hostB", 2)))) + + // The stage is available via its local completed-partition set... + assert(producerStage.isAvailable, + "pipelined producer should be available after its tasks finish") + assert(producerStage.isPipelined) + // ...but the pipelined shuffle has NO map outputs registered in the MapOutputTracker. + assert(mapOutputTracker.getNumAvailableOutputs(shuffleId) === 0, + "a pipelined shuffle must not register durable map outputs in the MapOutputTracker") + + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: numAvailableOutputs/findMissingPartitions reflect the exact " + + "incomplete partition set") { + // Pin the availability CONTRACT for a partially-complete pipelined producer, not just the + // all-done / none-done endpoints: numAvailableOutputs must be the count of completed partitions + // and findMissingPartitions must return the exact ids still missing (identity, not just size). + // Without this, a plausible refactor (revert the findMissingPartitions isPipelined branch to + // the + // tracker, use a bare counter, drop/take by size, or invert the filter) would silently break + // the + // contract yet pass every full-completion test. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + val producerTs = taskSets.head + val producerStage = + scheduler.stageIdToStage(producerTs.stageId).asInstanceOf[ShuffleMapStage] + + // Complete ONLY partition 1, leaving partition 0 missing. + runEvent(makeCompletionEvent(producerTs.tasks(1), Success, makeMapStatus("hostB", 2))) + assert(producerStage.numAvailableOutputs === 1) + assert(!producerStage.isAvailable) + assert(producerStage.findMissingPartitions() === Seq(0), + "findMissingPartitions must name the exact missing partition, not just a right-sized set") + + // Complete partition 0; now nothing is missing. + runEvent(makeCompletionEvent(producerTs.tasks(0), Success, makeMapStatus("hostA", 2))) + assert(producerStage.numAvailableOutputs === 2) + assert(producerStage.isAvailable) + assert(producerStage.findMissingPartitions() === Seq.empty) + + // Drain the consumer cleanly. + val consumerTs = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + complete(consumerTs, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a post-executor-loss straggler success must not resubmit the producer " + + "(bogus-epoch race)") { + // A pipelined producer's map task can succeed on an executor whose loss is already recorded + // (its StatusUpdate raced the executor-loss event). That completion hits the "possibly bogus + // epoch" branch, which for a regular shuffle simply ignores it (a healthy reattempt will + // re-register the output). But the same branch also runs `pendingPartitions -= partitionId` + // first, so if it is the last pending partition and we do NOT record it in the pipelined + // completed set, the stage looks "done but not available" and processShuffleMapStageCompletion + // resubmits the transient producer -- the exact streaming-writer hang. A pipelined stage must + // record the partition as completed even on the bogus-epoch path (its output is monotonic and + // MapOutputTracker's executor-loss stripping does not apply to it). + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + val producerTs = taskSets.head + val producerStage = + scheduler.stageIdToStage(producerTs.stageId).asInstanceOf[ShuffleMapStage] + + // Partition 0 completes normally on executor "hostA-exec" (recorded). Note makeMapStatus(host) + // builds executorId = host + "-exec", so pass host "hostA" to get executorId "hostA-exec". + runEvent(makeCompletionEvent(producerTs.tasks(0), Success, makeMapStatus("hostA", 2))) + assert(producerStage.numAvailableOutputs === 1) + val taskSetsBefore = taskSets.size + + // Executor "hostA-exec" is lost: this records executorFailureEpoch("hostA-exec") = current + // epoch. + runEvent(ExecutorLost("hostA-exec", ExecutorExited(-100, false, "Container marked as failed"))) + + // Now a delayed straggler Success for partition 1 arrives FROM the lost executor "hostA-exec". + // Its task epoch is the default (-1) <= the recorded failure epoch, so it is treated as + // possibly-bogus. It must still be recorded for a pipelined stage, so the producer becomes + // available and is NOT resubmitted. + runEvent(makeCompletionEvent(producerTs.tasks(1), Success, makeMapStatus("hostA", 2))) + + assert(producerStage.isAvailable, + "a pipelined producer must be available after all partitions succeed, even via a " + + "bogus-epoch " + + "straggler") + assert(producerStage.findMissingPartitions() === Seq.empty) + assert(taskSets.size === taskSetsBefore, + "the pipelined producer must NOT be resubmitted by a post-loss straggler success") + + // Consumer drains normally; no hang. + val consumerTs = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + complete(consumerTs, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a producer success with ignoreOldTaskAttempts set is still recorded " + + "(no resubmit)") { + // Sibling of the bogus-epoch case, for the OTHER guard that gates a regular shuffle's + // recording: + // ignoreOldTaskAttempts (set when a stage is rolled back, e.g. as a succeeding stage of an + // indeterminate ancestor). A pipelined producer's completed set is monotonic and never rolled + // back, so its success must be recorded even when ignoreOldTaskAttempts is true -- otherwise + // the + // last partition is dropped (pendingPartitions decremented, not recorded) -> "done but not + // available" -> processShuffleMapStageCompletion resubmits the transient producer. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + val producerTs = taskSets.head + val producerStage = + scheduler.stageIdToStage(producerTs.stageId).asInstanceOf[ShuffleMapStage] + + runEvent(makeCompletionEvent(producerTs.tasks(0), Success, makeMapStatus("hostA", 2))) + assert(producerStage.numAvailableOutputs === 1) + val taskSetsBefore = taskSets.size + + // Force ignoreOldTaskAttempts=true for the next completion: maxAttemptIdToIgnore >= the task's + // stageAttemptId (0). For a regular stage this would drop the completion; a pipelined stage + // must + // still record it. + producerStage.maxAttemptIdToIgnore = Some(0) + runEvent(makeCompletionEvent(producerTs.tasks(1), Success, makeMapStatus("hostB", 2))) + + assert(producerStage.isAvailable, + "a pipelined producer success must be recorded even when ignoreOldTaskAttempts is set") + assert(producerStage.findMissingPartitions() === Seq.empty) + assert(taskSets.size === taskSetsBefore, + "the pipelined producer must NOT be resubmitted when a success arrives under " + + "ignoreOldTaskAttempts") + + val consumerTs = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + complete(consumerTs, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: losing an executor does NOT flip a completed producer to unavailable " + + "or resubmit it") { + // A completed, consumed pipelined producer whose executor is lost must NOT be + // resubmitted (which would hang the streaming writer in awaitTerminationAcks). Because + // pipelined + // availability is tracked on the stage (monotonic) and not the MapOutputTracker, executor loss + // cannot flip isAvailable, so the producer is not resubmitted. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + val producerStageId = taskSets.head.stageId + val producerStage = + scheduler.stageIdToStage(producerStageId).asInstanceOf[ShuffleMapStage] + // Producer completes on hostA (both partitions). makeMapStatus("hostA") registers the output + // under executor id "hostA-exec" (makeBlockManagerId appends "-exec"), so the ExecutorLost + // below -- which loses executor id "hostA-exec" -- actually matches the registered output. + complete(taskSets.head, Seq( + (Success, makeMapStatus("hostA", 2)), + (Success, makeMapStatus("hostA", 2)))) + assert(producerStage.isAvailable) + val taskSetsAfterProducer = taskSets.size + + // Lose the executor that ran the producer ("hostA-exec"). For a regular shuffle this would + // strip the outputs and flip isAvailable -> resubmit; for a pipelined shuffle it must be inert. + runEvent(ExecutorLost("hostA-exec", ExecutorExited(-100, false, "Container marked as failed"))) + + assert(producerStage.isAvailable, + "a completed pipelined producer must remain available after executor loss (no tracker strip)") + // Pin the underlying completed set directly (not just the derived isAvailable boolean): + // executor + // loss must not remove any completed partition, so nothing is missing. + assert(producerStage.findMissingPartitions() === Seq.empty, + "executor loss must not strip a pipelined producer's completed partitions (monotonic)") + assert(taskSets.size === taskSetsAfterProducer, + "the pipelined producer must NOT be resubmitted on executor loss") + + // The consumer completes normally; no hang, no extra producer attempt. + val consumerTaskSet = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: binding a live producer stage to a second concurrent job fails fast") { + // A pipelined shuffle is a once-through live stream with no retained output, so two concurrent + // jobs cannot share one producer stage. While job 0 is still active its producer + // stage stays cached in shuffleIdToMapStage bound to job 0; a second concurrent job that reuses + // the SAME PipelinedShuffleDependency would bind that live stage to a second jobId, which is + // the + // forbidden cross-job reuse. Fail fast rather than let job 1 attach to job 0's live stream. + // (A sequential re-run is NOT this case: after job 0 finishes, cleanup drops the stage from + // shuffleIdToMapStage, so a later job gets a fresh producer bound to only its own job -- + // exactly how each streaming micro-batch reruns its producer.) + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + + // Job 0: submit but leave it active (do not complete the producer or the result stage), so the + // producer stage remains cached in shuffleIdToMapStage bound to job 0. + val firstConsumer = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(firstConsumer, Array(0, 1)) + val producerStage = + scheduler.stageIdToStage(taskSets.head.stageId).asInstanceOf[ShuffleMapStage] + assert(producerStage.jobIds === Set(0)) + val shuffleStagesBeforeJob1 = scheduler.shuffleIdToMapStage.size + val stagesBeforeJob1 = scheduler.stageIdToStage.size + + // Job 1: a second concurrent job reusing the SAME pipelined dependency -> cross-job reuse of a + // live producer stage -> fail fast. + val secondConsumer = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = {} + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(secondConsumer, Array(0, 1), listener = failListener) + assert(failure.get() != null, + "binding a live pipelined producer to a second concurrent job must fail the job") + assert(failure.get().getMessage.contains("PIPELINED_SHUFFLE_CROSS_JOB_REUSE") || + failure.get().getMessage.contains("reused across jobs"), + s"expected a cross-job-reuse error, got: ${failure.get().getMessage}") + // Job 0's producer stage was never bound to job 1, and the failed job 1 left no scheduler + // residue (the throw happened during stage creation, before any job-1 state was registered). + assert(producerStage.jobIds === Set(0)) + assert(scheduler.shuffleIdToMapStage.size === shuffleStagesBeforeJob1, + "the failed cross-job submission must not add a shuffle->stage mapping") + assert(scheduler.stageIdToStage.size === stagesBeforeJob1, + "the failed cross-job submission must not leave orphaned stages") + + // Job 0 can still complete normally -- the rejected job 1 did not corrupt its shared producer. + completeShuffleMapStageSuccessfully(producerStage.id, 0, 2) + val job0Consumer = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq firstConsumer + }.get + complete(job0Consumer, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: sequential re-run of the same producer is sound (fresh stage per job)") { + // Contrast with the concurrent case above: after a job using a pipelined dependency finishes, + // cleanup removes its producer from shuffleIdToMapStage, so re-submitting a job on the SAME + // dependency creates a fresh producer stage bound to only the new job. This is sound (it is how + // a streaming micro-batch reruns) and must NOT trip the cross-job-reuse fail-fast. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + + val consumer0 = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumer0, Array(0, 1)) + complete(taskSets(0), Seq( + (Success, makeMapStatus("hostA", 2)), + (Success, makeMapStatus("hostB", 2)))) + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + results.clear() + + // Second, sequential job on the same dependency: a fresh producer stage, no fail-fast, no hang. + // taskSets is a growing buffer across both jobs, so the second job's task sets are at index 2+. + val consumer1 = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumer1, Array(0, 1)) + complete(taskSets(2), Seq( + (Success, makeMapStatus("hostA", 2)), + (Success, makeMapStatus("hostB", 2)))) + complete(taskSets(3), Seq((Success, 4), (Success, 5))) + assert(results === Map(0 -> 4, 1 -> 5)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a FetchFailed on a group member fails the group, not a single-stage " + + "resubmit") { + // A FetchFailed must fail a pipelined (streaming) query promptly rather than trigger the + // base scheduler's single-stage resubmit -> serial recompute -> deadlock. The transient + // pipelined shuffle cannot be re-read, and members are co-scheduled, so a lone-stage resubmit + // is never valid. Any member's FetchFailed must abort the whole group (-> job abort -> + // the caller reruns the batch). Note the base TaskSetManager does NOT count a FetchFailed (it + // marks the task successful and zombies the set), so the group-atomic maxTaskFailures=1 lever + // does not apply to FetchFailed; the routing must be enforced in the DAGScheduler. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + val producerTs = taskSets.head + val producerStage = + scheduler.stageIdToStage(producerTs.stageId).asInstanceOf[ShuffleMapStage] + + // Producer completes (its outputs are tracked locally on the pipelined stage). + complete(producerTs, Seq( + (Success, makeMapStatus("hostA", 2)), + (Success, makeMapStatus("hostB", 2)))) + val consumerTs = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + val taskSetsBeforeFetchFailure = taskSets.size + + // A consumer task hits a FetchFailed reading the pipelined shuffle. + runEvent(makeCompletionEvent( + consumerTs.tasks(0), + FetchFailed(makeBlockManagerId("hostA"), pipelinedDep.shuffleId, 0L, 0, 0, "ignored"), + null)) + + // The job must be failed (group-atomic), and no single stage may be resubmitted: no new task + // set is created, and the scheduler is not left waiting to recompute the producer in isolation. + // scheduleResubmit posts ResubmitFailedStages on a timer, so drive any pending resubmit and + // confirm nothing new is launched. + scheduler.resubmitFailedStages() + assert(failure != null, "a FetchFailed on a pipelined group member must fail the job") + assert(taskSets.size === taskSetsBeforeFetchFailure, + "a pipelined group member's FetchFailed must NOT resubmit a single stage") + assert(!scheduler.runningStages.exists(_.isInstanceOf[ShuffleMapStage]), + "the pipelined producer must not be left running/resubmitted after the group fails") + sc.listenerBus.waitUntilEmpty() + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a FetchFailed reaching the base path aborts the group instead of " + + "throwing or resubmitting the transient producer") { + // A pipelined-group FetchFailed is normally intercepted by the group-atomic branch, whose guard + // is keyed on the failed stage's active job carrying hasPipelinedDependency. If that flag is + // ever not set/propagated for a job reading a pipelined shuffle, the FetchFailed falls through + // to the base path. That base path is invalid for a pipelined shuffle in two ways: (a) its + // MapOutputTracker invalidation would throw ShuffleStatusNotFoundException (a pipelined shuffle + // is never registered there); (b) its resubmit branch would enqueue a lone-stage resubmit of + // the transient producer, which cannot be re-read and would deadlock the group. The base path + // must instead abort the group. We reproduce the guard-miss by clearing the flag on the live + // active job, then driving the FetchFailed. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + val producerTs = taskSets.head + complete(producerTs, Seq( + (Success, makeMapStatus("hostA", 2)), + (Success, makeMapStatus("hostB", 2)))) + val consumerTs = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + val taskSetsBeforeFetchFailure = taskSets.size + + // Force the group-atomic guard to miss: clear hasPipelinedDependency on the live active job so + // the FetchFailed reaches the base path with a pipelined mapStage. + scheduler.activeJobs.foreach(_.hasPipelinedDependency = false) + assert(scheduler.shuffleIdToMapStage(pipelinedDep.shuffleId).isPipelined, + "the fetched shuffle must be the pipelined producer's") + + runEvent(makeCompletionEvent( + consumerTs.tasks(0), + FetchFailed(makeBlockManagerId("hostA"), pipelinedDep.shuffleId, 0L, 0, 0, "ignored"), + null)) + + // (a) The throwing MapOutputTracker invalidation was skipped for the pipelined id. + verify(mapOutputTracker, never()).unregisterAllMapAndMergeOutput(pipelinedDep.shuffleId) + verify(mapOutputTracker, never()).unregisterMapOutput( + org.mockito.ArgumentMatchers.eq(pipelinedDep.shuffleId), + org.mockito.ArgumentMatchers.anyInt(), + org.mockito.ArgumentMatchers.any()) + // (b) The group was aborted, NOT resubmitted: the job failed and no lone-stage resubmit of the + // transient producer was enqueued or launched. + scheduler.resubmitFailedStages() + assert(failure != null, "the group must be aborted (job failed) on the base-path FetchFailed") + assert(taskSets.size === taskSetsBeforeFetchFailure, + "a pipelined producer must NOT be resubmitted as a lone stage") + assert(!scheduler.runningStages.exists(_.isInstanceOf[ShuffleMapStage]), + "the transient pipelined producer must not be left running/resubmitted") + sc.listenerBus.waitUntilEmpty() + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a consumer result task throwing aborts the whole group") { + // Defense-in-depth for the group-atomic model at the DAGScheduler layer: a CONSUMER (result) + // task failing must tear down the whole group, not just its own stage. Result and map tasks + // share handleFailedTask / effectiveMaxTaskFailures=1, so the first consumer-task exception + // makes the TaskSetManager abort the set (delivered here as TaskSetFailed, exactly as a + // maxTaskFailures=1 abort surfaces to the DAGScheduler), which must abort the group and tear + // down the still-running producer -- the caller then reruns the batch. Complements the + // TaskSetManager-layer maxTaskFailures=1 tests and the producer-failure drop test. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + val consumerTs = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + val producerStage = scheduler.stageIdToStage(taskSets.head.stageId) + // The producer is STILL RUNNING (co-scheduled, not yet completed) when the consumer fails -- + // this is what gives the assertion teeth: group-atomic teardown must reach a running producer. + assert(scheduler.runningStages.contains(producerStage), + "the pipelined producer must be co-scheduled and running alongside the consumer") + + // The consumer's task set aborts on the first task exception (maxTaskFailures=1 for a pipelined + // member). This must fail the whole group and tear down the still-running producer. + failed(consumerTs, "consumer result task threw") + assert(failure != null, "a consumer task failure must fail the group's job") + assert(!scheduler.runningStages.contains(producerStage), + "the still-running pipelined producer must be torn down when the consumer fails the group") + sc.listenerBus.waitUntilEmpty() + assertDataStructuresEmpty() + } + + // ========================================================================================== + // Group-atomic rerun resets per-partition commit authorization + // ========================================================================================== + + test("pipelined shuffle: a group rerun resets per-partition commit authorization") { + // A pipelined group is atomic, so a failure reruns the WHOLE group -- including a result + // stage whose tasks already succeeded and committed. Those committed partitions are rerun and + // must be allowed to commit again. OutputCommitCoordinator permanently denies re-commit for a + // committed partition (a Success clears nothing; keyed by stage id). The rerun stays correct + // two ways, both asserted here: (b) the group teardown runs the committed result stage through + // markStageAsFinished -> stageEnd, clearing its committer state (so no stale authorization + // survives); and (a) the caller's rerun is a NEW job whose stages get FRESH stage ids, so the + // coordinator has no prior committer for them regardless. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = {} + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(consumerRdd, Array(0, 1), listener = failListener) + // Locate producer/consumer defensively by RDD identity (not task-set order). + val producerTaskSet = + taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq producerRdd).get + val consumerTaskSet = + taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd).get + val firstConsumerStageId = consumerTaskSet.stageId + + // Both stages register OutputCommitCoordinator state at submission (stageStart), so the + // coordinator holds per-stage commit-authorization state that a rerun must not inherit. (This + // asserts stage state exists to be reset; it does not claim a committer was authorized -- that + // needs a live canCommit RPC, which the mock backend does not drive.) + assert(!scheduler.outputCommitCoordinator.isEmpty, + "the coordinator must hold commit-authorization state for the submitted group's stages") + + // The producer now fails -> group-atomic failure -> the whole group is torn down and the job + // fails; the caller will rerun the batch as a new job. + failed(producerTaskSet, "producer blew up") + assert(failure.get() != null, "the group must fail atomically when the producer fails") + // Teardown ran the group's stages through markStageAsFinished -> stageEnd, clearing their + // coordinator state: the commit-authorization state does not survive the failed attempt. + // (isEmpty here proves stageEnd reached every stage the teardown covered.) + assert(scheduler.outputCommitCoordinator.isEmpty, + "group teardown must reset per-partition commit authorization; none may survive") + assertDataStructuresEmpty() + + // The caller reruns the batch as a NEW job on the same dependency. Its stages get fresh ids, so + // the coordinator has no prior committer, and the rerun's partition 0 can commit again. Only + // the + // task sets submitted from here on belong to the rerun (earlier ones' stages were cleaned up, + // so + // look them up defensively). + val taskSetsBeforeRerun = taskSets.size + val rerunConsumer = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(rerunConsumer, Array(0, 1)) + val rerunTaskSets = taskSets.drop(taskSetsBeforeRerun) + val rerunConsumerTaskSet = rerunTaskSets.find { ts => + scheduler.stageIdToStage.get(ts.stageId).exists(_.rdd eq rerunConsumer) + }.get + val rerunProducerTaskSet = rerunTaskSets.find { ts => + scheduler.stageIdToStage.get(ts.stageId).exists(_.rdd eq producerRdd) + }.get + assert(rerunConsumerTaskSet.stageId != firstConsumerStageId, + "the rerun's result stage must get a fresh stage id (fresh coordinator state)") + complete(rerunProducerTaskSet, Seq( + (Success, makeMapStatus("hostA", 2)), + (Success, makeMapStatus("hostB", 2)))) + complete(rerunConsumerTaskSet, Seq((Success, 7), (Success, 8))) + assert(results === Map(0 -> 7, 1 -> 8), "the rerun must complete, re-committing its partitions") + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a buffered consumer success is dropped on group abort, and its " + + "TaskEnd is flushed so the stage is not leaked as running") { + // A consumer's successful task has its whole completion event buffered while its producer runs + // (coarse model) -- no TaskEnd yet. If the group is then torn down before the producer finishes + // (here: the consumer's OTHER task hits a FetchFailed, which aborts the group), the buffered + // success's result must NOT be applied, but its TaskEnd IS flushed on the drop path. A listener + // that tracks active tasks (e.g. AppStatusListener, which removes a stage once activeTasks hits + // 0) therefore does not leak the consumer stage as perpetually running: the success's TaskStart + // was delivered, and its matching TaskEnd arrives at teardown even though the result is + // dropped. + val endedTaskIds = new java.util.concurrent.ConcurrentHashMap[Long, Boolean]() + val recordingListener = new SparkListener { + override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = + endedTaskIds.put(taskEnd.taskInfo.taskId, true) + } + sc.addSparkListener(recordingListener) + try { + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = results.put(index, result) + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(consumerRdd, Array(0, 1), listener = failListener) + // Identify the co-scheduled consumer (ResultStage over consumerRdd) and its still-running + // pipelined producer by RDD identity rather than task-set order. + val consumerTaskSet = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + assert(scheduler.dependentStageMap.size === 1, + "the consumer must be co-scheduled with a running producer (a deferral must exist)") + + // Consumer partition 0 succeeds. Its whole completion event is buffered (no job result and no + // TaskEnd yet). Give it a known taskId so we can assert the drop path flushes it later. + val bufferedTaskId = 7007L + runEvent(makeCompletionEvent(consumerTaskSet.tasks(0), Success, 42, + taskInfo = createFakeTaskInfoWithId(bufferedTaskId))) + sc.listenerBus.waitUntilEmpty() + assert(!endedTaskIds.containsKey(bufferedTaskId), + "the buffered consumer success must not emit its TaskEnd while the producer still runs") + assert(results.isEmpty, "the consumer's job result must still be deferred (producer running)") + assert(scheduler.dependentStageMap.get(scheduler.stageIdToStage(consumerTaskSet.stageId)) + .exists(_.delayedTaskCompletionEvents.nonEmpty), + "the consumer's completion event must be buffered while its producer runs") + + // The consumer's OTHER task now hits a FetchFailed. For a pipelined group member this aborts + // the whole group (group-atomic failure) WITHOUT the producer ever finishing. The buffered + // success is dropped (its result must not be applied), but its TaskEnd is flushed on the drop + // path so active-task-tracking listeners see the task finish. + runEvent(makeCompletionEvent( + consumerTaskSet.tasks(1), + FetchFailed(makeBlockManagerId("hostA"), pipelinedDep.shuffleId, 0L, 0, 0, "ignored"), + null)) + scheduler.resubmitFailedStages() + assert(failure.get() != null, "the job must fail") + + sc.listenerBus.waitUntilEmpty() + assert(results.isEmpty, + "the deferred consumer success must be dropped on abort, not applied as a result") + assert(endedTaskIds.containsKey(bufferedTaskId), + "the buffered consumer success's TaskEnd must be flushed on the drop path, so the stage " + + "is not leaked as perpetually running") + assertDataStructuresEmpty() + } finally { + sc.removeSparkListener(recordingListener) + } + } + + test("pipelined shuffle: a barrier producer stage is rejected") { + // A barrier stage exposes output only after a global sync, incompatible with incremental reads. + val producerRdd = new MyRDD(sc, 2, Nil).barrier().mapPartitions(iter => iter) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + assertPipelinedUnsupported(submitAndCaptureFailure(consumerRdd, Array(0, 1)), "barrier") + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a statically-indeterminate producer is rejected") { + // Indeterminate output's recovery is stage rollback-and-recompute, which a group never performs + // (so it never applies); reject rather than carry dead machinery. + val producerRdd = new MyRDD(sc, 2, Nil, indeterminate = true) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + assertPipelinedUnsupported(submitAndCaptureFailure(consumerRdd, Array(0, 1)), "indeterminate") + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a reliable RDD checkpoint in a PRODUCER's chain is rejected") { + // A reliable checkpoint writes a durable, lineage-truncated snapshot -> reintroduces cross-time + // reuse of a transient edge and needs a post-success recompute of the vanished input. Rejected + // by walking the producer's within-stage chain for a ReliableRDDCheckpointData. Keyed on + // checkpointData (not isCheckpointed): the write has not happened yet at group-creation time. + withTempDir { dir => + sc.setCheckpointDir(dir.getCanonicalPath) + val checkpointableRdd = new MyCheckpointRDD(sc, 2, Nil) + checkpointableRdd.checkpoint() // sets checkpointData to a ReliableRDDCheckpointData + assert(checkpointableRdd.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]])) + val pipelinedDep = new PipelinedShuffleDependency(checkpointableRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + assertPipelinedUnsupported( + submitAndCaptureFailure(consumerRdd, Array(0, 1)), "reliable RDD checkpoint") + assertDataStructuresEmpty() + } + } + + test("pipelined shuffle: a reliable RDD checkpoint in a CONSUMER's chain is rejected") { + // The rejection must cover a consumer member's chain too, not just the producer's: a consumer's + // transient input IS the pipelined shuffle, so a reliable checkpoint there would re-read the + // vanished stream on recompute. Purely all-pipelined shape (no regular shuffle -- those are + // rejected separately): producer --pipelined--> consumer(checkpointed, result). The consumer + // reads the pipelined shuffle and its within-stage chain carries the reliable checkpoint. + withTempDir { dir => + sc.setCheckpointDir(dir.getCanonicalPath) + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + // The consumer RDD reads the pipelined shuffle AND is reliably checkpointed. + val consumerRdd = new MyCheckpointRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + consumerRdd.checkpoint() + assert(consumerRdd.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]])) + assertPipelinedUnsupported( + submitAndCaptureFailure(consumerRdd, Array(0, 1)), "reliable RDD checkpoint") + assertDataStructuresEmpty() + } + } + + test("pipelined shuffle: a reliable checkpoint DOWNSTREAM in the consumer stage (not on the " + + "reading RDD) is still rejected") { + // Coverage for a checkpoint that is NOT on the RDD reading the pipelined shuffle, but on a + // narrow-dep child of it WITHIN the same consumer stage. Purely all-pipelined shape: + // producer --pipelined--> reads --(narrow)--> checkpointed (result). + // `reads` and `checkpointed` are one stage. Rooting the check at the pipelined-reading RDD and + // walking parents would MISS this (the checkpoint is downstream of the read); the check must + // instead recognize that `checkpointed`'s own within-stage chain reads a pipelined shuffle. + withTempDir { dir => + sc.setCheckpointDir(dir.getCanonicalPath) + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val readsRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + // A narrow-dep child of the reading RDD, in the SAME stage, that is reliably checkpointed, + // and is the result RDD (no regular shuffle after -- that would be separately rejected). + val checkpointedRdd = new MyCheckpointRDD(sc, 2, List(new OneToOneDependency(readsRdd))) + checkpointedRdd.checkpoint() + assert(checkpointedRdd.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]])) + assertPipelinedUnsupported( + submitAndCaptureFailure(checkpointedRdd, Array(0, 1)), "reliable RDD checkpoint") + assertDataStructuresEmpty() + } + } + + test("pipelined shuffle: a producer feeding more than one consumer (fan-out) is rejected") { + // 1:N fan-out is deferred to a later version and rejected up front here. Fan-out is + // detected at the RDD level -- two DISTINCT RDDs listing the same pipelined shuffle as a + // dependency -- so it is expressible in an all-pipelined job without any regular shuffle: two + // consumer RDDs both read the same pipelined producer, unioned by a narrow dependency into the + // result. checkPipelinedGroupsSupportedInRDDGraph counts 2 distinct consumers for the shuffle. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerA = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + val consumerB = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + val union = new MyRDD(sc, 2, + List(new OneToOneDependency(consumerA), new OneToOneDependency(consumerB)), + tracker = mapOutputTracker) + assertPipelinedUnsupported( + submitAndCaptureFailure(union, Array(0, 1)), "more than one consumer") + assertDataStructuresEmpty() + } + + // Resource-profile rejection. The gang slot check measures capacity against the DEFAULT + // resource profile, so the whole group must run on the default profile; any member with an + // explicit non-default profile is rejected. The three shapes below must all be rejected; the + // uniform-non-default and non-default-plus-default cases in particular would each pass a + // "more than one DISTINCT profile" check yet still be admitted against the wrong (default) pool. + private def rpWithCores(cores: Int, cpus: Int): ResourceProfile = + new ResourceProfileBuilder() + .require(new ExecutorResourceRequests().cores(cores)) + .require(new TaskResourceRequests().cpus(cpus)).build() + + test("pipelined shuffle: a group with two distinct non-default resource profiles is rejected") { + val producerRdd = new MyRDD(sc, 2, Nil).withResources(rpWithCores(4, 1)) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + .withResources(rpWithCores(8, 2)) + assertPipelinedUnsupported( + submitAndCaptureFailure(consumerRdd, Array(0, 1)), "non-default resource profile") + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a group uniformly on one non-default resource profile is rejected") { + // Both members share ONE non-default profile -- a "distinct profile count" of 1, which a + // more-than-one-distinct-profile check would wrongly admit, then measure against the default + // profile's capacity (the wrong pool). Must be rejected: the whole group is off the default RP. + val rp = rpWithCores(4, 1) + val producerRdd = new MyRDD(sc, 2, Nil).withResources(rp) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + .withResources(rp) + assertPipelinedUnsupported( + submitAndCaptureFailure(consumerRdd, Array(0, 1)), "non-default resource profile") + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a group mixing a non-default profile with the default is rejected") { + // The producer carries an explicit non-default profile; the consumer is left on the default. + // The set of EXPLICIT profiles is again size 1, so a distinct-explicit-profile check would miss + // it, but the group genuinely spans the non-default and default profiles. Must be rejected. + val producerRdd = new MyRDD(sc, 2, Nil).withResources(rpWithCores(4, 1)) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + assertPipelinedUnsupported( + submitAndCaptureFailure(consumerRdd, Array(0, 1)), "non-default resource profile") + assertDataStructuresEmpty() + } + + test("regular shuffle idioms are NOT rejected (inertness of the pipelined fail-fast)") { + // The pipelined fail-fast checks (indeterminate producer, reliable-checkpoint-in-chain) must be + // inert for a job with NO pipelined dependency: a regular shuffle whose producer is BOTH + // indeterminate AND reliably checkpointed -- the two idioms most likely to false-fire -- must + // run exactly as before. (Actually exercise both, not just the flag.) + withTempDir { dir => + sc.setCheckpointDir(dir.getCanonicalPath) + val producerRdd = new MyCheckpointRDD(sc, 2, Nil, indeterminate = true) + producerRdd.checkpoint() + assert(producerRdd.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]])) + assert(producerRdd.outputDeterministicLevel == DeterministicLevel.INDETERMINATE) + val regularDep = new ShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(regularDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + // The job proceeds normally (producer stage submitted), i.e. NOT failed by a pipelined check. + assert(failure === null, "a regular shuffle must not be rejected by the pipelined fail-fast") + completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2) + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + } + + test("regular job on a non-default resource profile is NOT rejected (RP check is pipelined)") { + // The resource-profile rejection is not keyed on a pipelined dependency, so it must run ONLY + // for a job that has one (handleJobSubmitted gates checkPipelinedGroupsSupportedInRDDGraph on + // hasPipelined). A perfectly ordinary job that merely attaches a non-default profile via + // RDD.withResources -- a GA stage-level-scheduling feature -- has NO pipelined dependency and + // must run untouched. Without the gate the whole graph walk fires and rejects it with + // PIPELINED_SHUFFLE_UNSUPPORTED, a regression on plain withResources jobs. + val rdd = new MyRDD(sc, 2, Nil).withResources(rpWithCores(4, 2)) + // A non-default profile drives submitMissingTasks through addPySparkConfigsToProperties, which + // needs a non-null Properties; pass one (the default submit() overload leaves it null). + submit(rdd, Array(0, 1), properties = new Properties()) + assert(failure === null, + "a regular job with a non-default resource profile must not be rejected by the pipelined " + + "fail-fast") + assert(taskSets.nonEmpty, "the job must proceed to task submission, not be failed up front") + complete(taskSets(0), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + private def submitAndCaptureFailure(finalRdd: RDD[_], partitions: Array[Int]): Exception = { + val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = {} + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(finalRdd, partitions, listener = failListener) + failure.get() + } + + private def assertPipelinedUnsupported(failure: Exception, reasonSubstring: String): Unit = { + assert(failure != null, "the job must fail fast on the unsupported pipelined idiom") + val msg = failure.getMessage + assert(msg.contains("PIPELINED_SHUFFLE_UNSUPPORTED") || msg.contains("unsupported feature"), + s"expected a PIPELINED_SHUFFLE_UNSUPPORTED error, got: $msg") + assert(msg.contains(reasonSubstring), + s"expected reason to mention '$reasonSubstring', got: $msg") + } + } class DAGSchedulerAbortStageOffSuite extends DAGSchedulerSuite { diff --git a/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala index 28d774eb27f7e..4c7e8efc4eaba 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala @@ -3046,6 +3046,204 @@ class TaskSetManagerSuite "ExceptionFailure must count towards task failures (contrast)") } + // ========================================================================================== + // Pipelined-group member: group-atomic failure via job-abort + // ========================================================================================== + + /** A single-task TaskSet marked as a pipelined-group member. */ + private def pipelinedTaskSet(): TaskSet = { + val tasks = Array.tabulate[Task[_]](1)(i => new FakeTask(0, i, Nil)) + new TaskSet(tasks, stageId = 0, stageAttemptId = 0, priority = 0, null, + ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID, None, isPipelined = true) + } + + test("pipelined task set aborts after a single task failure (maxTaskFailures capped at 1)") { + sc = new SparkContext("local", "test") + sched = new FakeTaskScheduler(sc, ("exec1", "host1")) + val taskSet = pipelinedTaskSet() + val clock = new ManualClock + clock.advance(1) + // Even though MAX_TASK_FAILURES (4) is passed, a pipelined task set caps attempts at 1. + val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES, clock = clock) + + val offerResult = manager.resourceOffer("exec1", "host1", ANY)._1 + assert(offerResult.isDefined) + // A single ordinary task failure must abort the whole task set (group -> job), no retry. + manager.handleFailedTask(offerResult.get.taskId, TaskState.FINISHED, TaskResultLost) + assert(sched.taskSetsFailed.contains(taskSet.id), + "a pipelined task set must abort after the first task failure") + } + + test("pipelined task set counts an executor-loss failure that would normally be uncounted") { + sc = new SparkContext("local", "test") + sched = new FakeTaskScheduler(sc, ("exec1", "host1")) + val taskSet = pipelinedTaskSet() + val clock = new ManualClock + clock.advance(1) + val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES, clock = clock) + + val offerResult = manager.resourceOffer("exec1", "host1", ANY)._1 + assert(offerResult.isDefined) + // ExecutorLostFailure with exitCausedByApp=false is normally NOT counted toward task failures, + // so a non-pipelined task set would not abort. For a pipelined task set it counts, so the + // single failure aborts the group. + manager.handleFailedTask(offerResult.get.taskId, TaskState.FINISHED, + ExecutorLostFailure("exec1", exitCausedByApp = false, reason = None)) + assert(sched.taskSetsFailed.contains(taskSet.id), + "a pipelined task set must count executor loss and abort") + } + + test("pipelined task set does NOT abort on a benign TaskKilled (losing duplicate attempt)") { + // A pipelined set with 2 tasks. One task gets a duplicate attempt; when one wins, the loser is + // killed with TaskKilled. That benign kill must NOT count as a failure and must NOT abort the + // group (the task actually succeeded). Regression for the speculative-loser abort bug. + sc = new SparkContext("local", "test") + sched = new FakeTaskScheduler(sc, ("exec1", "host1")) + val tasks = Array.tabulate[Task[_]](2)(i => new FakeTask(0, i, Nil)) + val taskSet = new TaskSet(tasks, stageId = 0, stageAttemptId = 0, priority = 0, null, + ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID, None, isPipelined = true) + val clock = new ManualClock + clock.advance(1) + val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES, clock = clock) + + // Launch task index 0 and fail it with a benign TaskKilled (as if a duplicate attempt won). + val offer = manager.resourceOffer("exec1", "host1", ANY)._1 + assert(offer.isDefined && offer.get.index === 0) + manager.handleFailedTask(offer.get.taskId, TaskState.KILLED, + TaskKilled("another attempt succeeded")) + // The group must NOT be aborted: TaskKilled is benign even for a pipelined set. + assert(!sched.taskSetsFailed.contains(taskSet.id), + "a benign TaskKilled must not abort a pipelined group") + } + + test("pipelined task set does NOT abort on TaskCommitDenied (speculation commit race)") { + sc = new SparkContext("local", "test") + sched = new FakeTaskScheduler(sc, ("exec1", "host1")) + val tasks = Array.tabulate[Task[_]](2)(i => new FakeTask(0, i, Nil)) + val taskSet = new TaskSet(tasks, stageId = 0, stageAttemptId = 0, priority = 0, null, + ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID, None, isPipelined = true) + val clock = new ManualClock + clock.advance(1) + val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES, clock = clock) + + val offer = manager.resourceOffer("exec1", "host1", ANY)._1 + assert(offer.isDefined) + manager.handleFailedTask(offer.get.taskId, TaskState.FINISHED, + TaskCommitDenied(jobID = 0, partitionID = 0, attemptNumber = 0)) + assert(!sched.taskSetsFailed.contains(taskSet.id), + "TaskCommitDenied must not abort a pipelined group") + } + + test("pipelined task set does NOT single-resubmit a completed map task on executor " + + "decommission") { + // The TaskSetManager.executorLost "Resubmitted" loop re-enqueues an already-successful + // ShuffleMapTask when its executor is lost and the map output looks gone -- for a decommission, + // it looks gone whenever MapOutputTracker.getMapOutputLocation returns None. A pipelined + // shuffle is NEVER registered in the MapOutputTracker, so getMapOutputLocation is always None + // and this loop would resubmit the lone producer task -- the exact streaming-writer hang. This + // loop bypasses handleFailedTask, so the group-atomic abort would never see it. The guard + // (!taskSet.isPipelined on maybeShuffleMapOutputLoss) must keep a pipelined set out of the + // loop. + // + // Scenario: a 2-task pipelined producer, one task succeeds on the to-be-lost executor while the + // other is still running (so the set is NOT a zombie and the loop is entered for a + // non-pipelined set). Decommission that executor. Assert NO Resubmitted is emitted for the + // completed task. + sc = new SparkContext("local", "test") + sched = new FakeTaskScheduler(sc, ("execA", "host1"), ("execB", "host2")) + sched.initialize(new FakeSchedulerBackend()) + + var resubmittedTasks = 0 + val dagScheduler = new FakeDAGScheduler(sc, sched) { + override def taskEnded( + task: Task[_], + reason: TaskEndReason, + result: Any, + accumUpdates: Seq[AccumulatorV2[_, _]], + metricPeaks: Array[Long], + taskInfo: TaskInfo): Unit = { + super.taskEnded(task, reason, result, accumUpdates, metricPeaks, taskInfo) + reason match { + case Resubmitted => resubmittedTasks += 1 + case _ => + } + } + } + sched.dagScheduler.stop() + sched.setDAGScheduler(dagScheduler) + + // Two real ShuffleMapTasks (so isShuffleMapTasks = true), marked as a pipelined-group member. + // Give each task a distinct preferred location so one lands on execA and the other on execB. + val prefs = Array( + Seq[TaskLocation](TaskLocation("host1", "execA")), + Seq[TaskLocation](TaskLocation("host2", "execB"))) + val tasks = Array.tabulate[Task[_]](2) { i => + new ShuffleMapTask(0, 0, null, new Partition { override def index: Int = i }, 1, + prefs(i), JobArtifactSet.getActiveOrDefault(sc), new Properties, null) + } + val taskSet = new TaskSet(tasks, stageId = 0, stageAttemptId = 0, priority = 0, null, + ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID, Some(0), isPipelined = true) + // A pipelined shuffle is deliberately NOT registered in the MapOutputTracker, so a decommission + // lookup of its output returns None (looks lost) -- the trigger for channel 3. + val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES) + + // Launch task 0 on execA (PROCESS_LOCAL) and task 1 on execB; complete task 0 so the set is not + // yet a zombie (task 1 still running). + val t0 = manager.resourceOffer("execA", "host1", TaskLocality.PROCESS_LOCAL)._1.get + val t1 = manager.resourceOffer("execB", "host2", TaskLocality.PROCESS_LOCAL)._1.get + assert(manager.runningTasks === 2) + assert(t0.index === 0 && t1.index === 1, "each task should land on its preferred executor") + val result0 = new DirectTaskResult[String]() { + override def value(resultSer: SerializerInstance): String = "" + } + manager.handleSuccessfulTask(t0.taskId, result0) + assert(!manager.isZombie, "the set must not be a zombie (task 1 still running)") + assert(manager.successful(t0.index)) + assert(resubmittedTasks === 0) + + // Decommission execA (which ran the completed task 0) and lose it. For a NON-pipelined set this + // would re-enqueue task 0 via Resubmitted (its output looks lost); for a pipelined set the + // guard must prevent that. + manager.executorDecommission("execA") + manager.executorLost("execA", "host1", ExecutorDecommission()) + + assert(resubmittedTasks === 0, + "a pipelined producer's completed task must NOT be single-resubmitted on decommission") + assert(manager.successful(t0.index), + "the completed task must remain successful (no Resubmitted re-enqueue)") + } + + test("non-pipelined task set is unaffected: executor loss is not counted, retries still apply") { + sc = new SparkContext("local", "test") + sched = new FakeTaskScheduler(sc, ("exec1", "host1")) + val taskSet = FakeTask.createTaskSet(1) // isPipelined defaults to false + val clock = new ManualClock + clock.advance(1) + val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES, clock = clock) + + // An uncounted executor-loss failure must NOT abort a normal task set. + val offer1 = manager.resourceOffer("exec1", "host1", ANY)._1 + assert(offer1.isDefined) + manager.handleFailedTask(offer1.get.taskId, TaskState.FINISHED, + ExecutorLostFailure("exec1", exitCausedByApp = false, reason = None)) + assert(!sched.taskSetsFailed.contains(taskSet.id), + "a normal task set must not count uncaused executor loss") + + // And a normal task set still retries a counted failure up to MAX_TASK_FAILURES before abort. + (1 to MAX_TASK_FAILURES).foreach { index => + val offer = manager.resourceOffer("exec1", "host1", ANY)._1 + assert(offer.isDefined, s"expected an offer on iteration $index") + manager.handleFailedTask(offer.get.taskId, TaskState.FINISHED, TaskResultLost) + if (index < MAX_TASK_FAILURES) { + assert(!sched.taskSetsFailed.contains(taskSet.id), + s"must not abort before $MAX_TASK_FAILURES failures (iteration $index)") + } else { + assert(sched.taskSetsFailed.contains(taskSet.id), + "must abort after MAX_TASK_FAILURES counted failures") + } + } + } + } class FakeLongTasks(stageId: Int, partitionId: Int) extends FakeTask(stageId, partitionId) { diff --git a/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala b/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala index da74803a10ebc..8e0f9e04c1bd8 100644 --- a/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala +++ b/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala @@ -240,6 +240,46 @@ class PipelinedShuffleRoutingSuite extends SparkFunSuite with LocalSparkContext assert(SparkEnv.get.streamingShuffleOutputTracker.isEmpty) } + test("createShuffleMapStage fails loud for a pipelined dependency when no streaming tracker") { + // A pipelined shuffle is served solely by the StreamingShuffleOutputTracker. With a non- + // streaming incremental manager configured, that tracker is absent (verified above), so a job + // whose stage graph contains a PipelinedShuffleDependency cannot be scheduled. + // createShuffleMapStage must fail loud rather than silently register the shuffle in no tracker + // (which would strand a consumer with no writer locations). This guards the decoupling's + // invariant that a pipelined dependency implies a streaming tracker. + // local[4]: the pipelined group's up-front slot admission (producer 2 + result 2 = 4) must pass + // so that stage creation is actually reached -- otherwise the job is rejected for slots first. + sc = new SparkContext("local[4]", "test", newConf()) + assert(SparkEnv.get.streamingShuffleOutputTracker.isEmpty) + val producer = sc.parallelize(1 to 4, 2).map(x => (x, x)) + val pipelined = new PipelinedShuffleDependency[Int, Int, Int](producer, new HashPartitioner(2)) + // A minimal reduce-side RDD whose single dependency is the pipelined shuffle, so submitting an + // action forces createShuffleMapStage for the pipelined producer. + val consumer = new RDD[(Int, Int)](sc, Seq(pipelined)) { + override def compute(split: Partition, ctx: TaskContext): Iterator[(Int, Int)] = + Iterator.empty + override protected def getPartitions: Array[Partition] = + Array.tabulate(2)(i => new Partition { override def index: Int = i }) + } + val ex = intercept[Exception] { + consumer.count() + } + assert(findCause[IllegalStateException](ex).exists( + _.getMessage.contains("requires a StreamingShuffleOutputTracker")), + s"expected a fail-loud IllegalStateException about the missing tracker, got: $ex") + + // The fail-loud throw happens BEFORE createShuffleMapStage mutates stageIdToStage / + // shuffleIdToMapStage, so it leaves no partial scheduler state behind. If it left a + // half-created stage cached in shuffleIdToMapStage, a re-submission would reuse that stale + // stage instead of re-throwing. Re-submitting the same job must therefore re-throw the error. + val ex2 = intercept[Exception] { + consumer.count() + } + assert(findCause[IllegalStateException](ex2).exists( + _.getMessage.contains("requires a StreamingShuffleOutputTracker")), + s"a re-submission must re-throw the fail-loud error (no leaked partial stage), got: $ex2") + } + test("spark.shuffle.manager.incremental resolves the same short aliases as the default manager") { // "sort" is a short alias the incremental slot must resolve (to SortShuffleManager) rather than // treat as a class name. SortShuffleManager is blocking, so it is rejected from the pipelined diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffset.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffset.scala index aa29128cda7e0..88f0de2555f70 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffset.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffset.scala @@ -44,7 +44,7 @@ object InsertSortForLimitAndOffset extends Rule[SparkPlan] { _, // Should not match AQE shuffle stage because we only target un-submitted stages which // we can still rewrite the query plan. - s @ ShuffleExchangeExec(SinglePartition, child, _, _), + s @ ShuffleExchangeExec(SinglePartition, child, _, _, _), _) if child.logicalLink.isDefined => extractOrderingAndPropagateOrderingColumns(child) match { case Some((ordering, newChild)) => diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEUtils.scala index 578e0acd80525..cc3ca90739090 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEUtils.scala @@ -30,7 +30,7 @@ object AQEUtils { // Project/Filter/LocalSort/CollectMetrics. // Note: we only care about `HashPartitioning` as `EnsureRequirements` can only optimize out // user-specified repartition with `HashPartitioning`. - case ShuffleExchangeExec(h: HashPartitioning, _, shuffleOrigin, _) + case ShuffleExchangeExec(h: HashPartitioning, _, shuffleOrigin, _, _) if shuffleOrigin == REPARTITION_BY_COL || shuffleOrigin == REPARTITION_BY_NUM => val numPartitions = if (shuffleOrigin == REPARTITION_BY_NUM) { Some(h.numPartitions) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala index c632b3d841e61..8ca344e511946 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala @@ -272,8 +272,8 @@ case class EnsureRequirements( } child match { - case ShuffleExchangeExec(_, c, so, ps) => - ShuffleExchangeExec(newPartitioning, c, so, ps) + case s: ShuffleExchangeExec => + s.copy(outputPartitioning = newPartitioning) case gpe: GroupPartitionsExec => ShuffleExchangeExec(newPartitioning, gpe.child) case _ => ShuffleExchangeExec(newPartitioning, child) } @@ -896,7 +896,7 @@ case class EnsureRequirements( def apply(plan: SparkPlan): SparkPlan = { val newPlan = plan.transformUp { - case operator @ ShuffleExchangeExec(upper: HashPartitioning, child, shuffleOrigin, _) + case operator @ ShuffleExchangeExec(upper: HashPartitioning, child, shuffleOrigin, _, _) if optimizeOutRepartition && (shuffleOrigin == REPARTITION_BY_COL || shuffleOrigin == REPARTITION_BY_NUM) => def hasSemanticEqualPartitioning(partitioning: Partitioning): Boolean = { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala index dd829e697df61..92893a332ba08 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala @@ -191,7 +191,8 @@ case class ShuffleExchangeExec( override val outputPartitioning: Partitioning, child: SparkPlan, shuffleOrigin: ShuffleOrigin = ENSURE_REQUIREMENTS, - advisoryPartitionSize: Option[Long] = None) + advisoryPartitionSize: Option[Long] = None, + pipelined: Boolean = false) extends ShuffleExchangeLike { private lazy val writeMetrics = @@ -252,7 +253,8 @@ case class ShuffleExchangeExec( child.output, outputPartitioning, serializer, - writeMetrics) + writeMetrics, + pipelined) metrics("numPartitions").set(dep.partitioner.numPartitions) val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY) SQLMetrics.postDriverMetricUpdates( @@ -346,7 +348,8 @@ object ShuffleExchangeExec { outputAttributes: Seq[Attribute], newPartitioning: Partitioning, serializer: Serializer, - writeMetrics: Map[String, SQLMetric]) + writeMetrics: Map[String, SQLMetric], + pipelined: Boolean = false) : ShuffleDependency[Int, InternalRow, InternalRow] = { val part: Partitioner = newPartitioning match { case RoundRobinPartitioning(numPartitions) => new HashPartitioner(numPartitions) @@ -537,15 +540,30 @@ object ShuffleExchangeExec { } } val dependency = - new ShuffleDependency[Int, InternalRow, InternalRow]( - rddWithPartitionIds, - new PartitionIdPassthrough(part.numPartitions), - serializer, - shuffleWriterProcessor = createShuffleWriteProcessor(writeMetrics), - rowBasedChecksums = UnsafeRowChecksum.createUnsafeRowChecksums(checksumSize), - _checksumMismatchFullRetryEnabled = SQLConf.get.shuffleChecksumMismatchFullRetryEnabled, - checksumMismatchQueryLevelRollbackEnabled = - SQLConf.get.shuffleChecksumMismatchQueryLevelRollbackEnabled) + if (pipelined) { + // A pipelined shuffle is transient and incrementally readable: the DAGScheduler + // co-schedules its producer and consumer stages instead of materializing the shuffle + // first. The PipelinedShuffleDependency type is the entire opt-in -- routing to the + // streaming shuffle manager and pipelined-group co-scheduling both follow from it. The + // checksum-mismatch retry knobs are intentionally not carried over: a transient shuffle is + // never recomputed, so PipelinedShuffleDependency does not expose them (they stay off). + new PipelinedShuffleDependency[Int, InternalRow, InternalRow]( + rddWithPartitionIds, + new PartitionIdPassthrough(part.numPartitions), + serializer, + shuffleWriterProcessor = createShuffleWriteProcessor(writeMetrics), + rowBasedChecksums = UnsafeRowChecksum.createUnsafeRowChecksums(checksumSize)) + } else { + new ShuffleDependency[Int, InternalRow, InternalRow]( + rddWithPartitionIds, + new PartitionIdPassthrough(part.numPartitions), + serializer, + shuffleWriterProcessor = createShuffleWriteProcessor(writeMetrics), + rowBasedChecksums = UnsafeRowChecksum.createUnsafeRowChecksums(checksumSize), + _checksumMismatchFullRetryEnabled = SQLConf.get.shuffleChecksumMismatchFullRetryEnabled, + checksumMismatchQueryLevelRollbackEnabled = + SQLConf.get.shuffleChecksumMismatchQueryLevelRollbackEnabled) + } dependency } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala index 0d2e4a6941a00..a037369e3addb 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala @@ -36,8 +36,9 @@ import org.apache.spark.sql.classic.SparkSession import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.execution.{CommandExecutionMode, LocalLimitExec, QueryExecution, SerializeFromObjectExec, SparkPlan, SparkPlanner, SparkStrategy => Strategy, UnaryExecNode} import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, MergingSessionsExec, ObjectHashAggregateExec, SortAggregateExec, UpdatingSessionsExec} +import org.apache.spark.sql.execution.datasources.v2.RealTimeStreamScanExec import org.apache.spark.sql.execution.datasources.v2.state.metadata.StateMetadataPartitionReader -import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike +import org.apache.spark.sql.execution.exchange.{ShuffleExchangeExec, ShuffleExchangeLike} import org.apache.spark.sql.execution.python.streaming.{FlatMapGroupsInPandasWithStateExec, TransformWithStateInPySparkExec} import org.apache.spark.sql.execution.streaming.{StreamingErrors, StreamingQueryPlanTraverseHelper} import org.apache.spark.sql.execution.streaming.checkpointing.{CheckpointFileManager, OffsetSeqMetadata, OffsetSeqMetadataBase} @@ -663,7 +664,44 @@ class IncrementalExecution( } } - override def preparations: Seq[Rule[SparkPlan]] = state +: super.preparations + /** + * For a Real-Time Mode batch, mark the shuffle exchanges as pipelined so the DAGScheduler + * co-schedules a stateful query's producer (source scan) and consumer (stateful operator) stages + * as one pipelined group -- records stream through a transient shuffle instead of the consumer + * waiting for the producer to fully materialize. The exchange carries the decision as a field + * (see ShuffleExchangeExec.pipelined); the PipelinedShuffleDependency it then builds is the whole + * opt-in -- routing to the streaming shuffle manager and pipelined-group scheduling both follow + * from that dependency type. + * + * Real-Time Mode is detected structurally by a RealTimeStreamScanExec leaf (there is no + * RTM-specific plan flag). Inert for a non-RTM batch, so the ordinary microbatch path is + * unchanged. + * + * Single-shuffle assumption: today an RTM plan has exactly one shuffle exchange (a second one is + * rejected by RealTimeModeAllowlist -- see the "repartition not allowed" test / SPARK-54237), so + * transformUp marks that one exchange and there is no exchange reuse to worry about. If RTM is + * ever widened to multi-shuffle plans (stateful joins, CTEs read twice, a union of the same keyed + * stream), transformUp would NOT reach a ReusedExchangeExec's wrapped `child` (it is a leaf whose + * child is a field, not a tree child), so a reused shuffle would stay regular while its in-tree + * twin becomes pipelined -- a mixed regular+pipelined job the scheduler rejects every batch. When + * that day comes, mark inside ReusedExchangeExec (or run this before exchange reuse, or disable + * reuse for RTM) and add a multi-shuffle test. + */ + object MarkPipelinedShuffleForRealTimeMode extends Rule[SparkPlan] { + override def apply(plan: SparkPlan): SparkPlan = { + val isRealTimeMode = plan.exists(_.isInstanceOf[RealTimeStreamScanExec]) + if (!isRealTimeMode) { + plan + } else { + plan.transformUp { + case s: ShuffleExchangeExec if !s.pipelined => s.copy(pipelined = true) + } + } + } + } + + override def preparations: Seq[Rule[SparkPlan]] = + state +: (super.preparations :+ MarkPipelinedShuffleForRealTimeMode) /** no need to try-catch again as this is already done once */ override def assertAnalyzed(): Unit = analyzed diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala index 7ae557797f36c..047ea3de4298d 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala @@ -55,7 +55,17 @@ object RealTimeModeAllowlist extends Logging { "org.apache.spark.sql.execution.datasources.v2.WriteToDataSourceV2Exec", "org.apache.spark.sql.execution.exchange.BroadcastExchangeExec", "org.apache.spark.sql.execution.exchange.ReusedExchangeExec", + // A pipelined (streaming) shuffle repartitions a stateful query's input by key. In Real-Time + // Mode the DAGScheduler co-schedules the shuffle's producer and consumer stages via a + // PipelinedShuffleDependency (see IncrementalExecution's pipelined-shuffle rule), so the + // exchange is a supported member of a pipelined group rather than a materialization barrier. + "org.apache.spark.sql.execution.exchange.ShuffleExchangeExec", "org.apache.spark.sql.execution.joins.BroadcastHashJoinExec", + // Streaming deduplication and the state-store access operators it plans into. These run in the + // pipelined-shuffle consumer stage, keyed by the same columns the shuffle repartitions on. + "org.apache.spark.sql.execution.streaming.operators.stateful.StateStoreRestoreExec", + "org.apache.spark.sql.execution.streaming.operators.stateful.StateStoreSaveExec", + "org.apache.spark.sql.execution.streaming.operators.stateful.StreamingDeduplicateExec", classOf[EventTimeWatermarkExec].getName ) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala index c1d1f13915a6e..136dc57fdfcfa 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala @@ -1360,7 +1360,7 @@ class DataFrameWindowFunctionsSuite extends SharedSparkSession def isShuffleExecByRequirement( plan: ShuffleExchangeExec, desiredClusterColumns: Seq[String]): Boolean = plan match { - case ShuffleExchangeExec(op: HashPartitioning, _, ENSURE_REQUIREMENTS, _) => + case ShuffleExchangeExec(op: HashPartitioning, _, ENSURE_REQUIREMENTS, _, _) => partitionExpressionsColumns(op.expressions) === desiredClusterColumns case _ => false } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala index 3b28cae31a134..096eee137b320 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala @@ -1942,7 +1942,7 @@ class DatasetSuite extends SharedSparkSession val agg = cp.groupBy($"id" % 2).agg(count($"id")) agg.queryExecution.executedPlan.collectFirst { - case ShuffleExchangeExec(_, _: RDDScanExec, _, _) => + case ShuffleExchangeExec(_, _: RDDScanExec, _, _, _) => case BroadcastExchangeExec(_, _: RDDScanExec) => }.foreach { _ => fail( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala index 71192e92af0d0..41b74c6d4ff21 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala @@ -733,10 +733,11 @@ class PlannerSuite extends SharedSparkSession with AdaptiveSparkPlanHelper { outputPlan match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, - ShuffleExchangeExec(HashPartitioning(leftPartitioningExpressions, _), _, _, _), _), + ShuffleExchangeExec(HashPartitioning(leftPartitioningExpressions, _), + _, _, _, _), _), SortExec(_, _, ShuffleExchangeExec(HashPartitioning(rightPartitioningExpressions, _), - _, _, _), _), _) => + _, _, _, _), _), _) => assert(leftKeys === smjExec.leftKeys) assert(rightKeys === smjExec.rightKeys) assert(leftKeys === leftPartitioningExpressions) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala index 17d00ec055e07..95085b71bc251 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala @@ -60,7 +60,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(smjExec1) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: PartitioningCollection, _, _), _), - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB)) assert(rightKeys === Seq(exprB, exprA)) case other => fail(other.toString) @@ -71,7 +71,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprB :: exprA :: Nil, Inner, None, plan2, plan1) EnsureRequirements.apply(smjExec2) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, _: PartitioningCollection, _, _), _), _) => assert(leftKeys === Seq(exprB, exprA)) assert(rightKeys === Seq(exprA, exprB)) @@ -84,7 +84,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprD :: exprC :: Nil, exprB :: exprA :: Nil, Inner, None, plan1, plan1) EnsureRequirements.apply(smjExec3) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, _: PartitioningCollection, _, _), _), _) => assert(leftKeys === Seq(exprC, exprD)) assert(rightKeys === Seq(exprA, exprB)) @@ -126,8 +126,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { ) EnsureRequirements.apply(smjExec2) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprC, exprB, exprD)) assert(rightKeys === Seq(exprD, exprA, exprC)) case other => fail(other.toString) @@ -145,7 +145,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprB :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec1) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), _) => assert(leftKeys === Seq(exprB, exprA)) assert(rightKeys === Seq(exprB, exprC)) @@ -159,7 +159,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprB :: Nil, Inner, None, plan1, plan3) EnsureRequirements.apply(smjExec2) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, _: PartitioningCollection, _, _), _), _) => assert(leftKeys === Seq(exprB, exprA)) assert(rightKeys === Seq(exprB, exprC)) @@ -173,7 +173,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(smjExec3) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: PartitioningCollection, _, _), _), - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprB, exprC)) assert(rightKeys === Seq(exprB, exprA)) case other => fail(other.toString) @@ -319,7 +319,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB)) assert(rightKeys === Seq(exprC, exprD)) assert(p.expressions == Seq(exprC)) @@ -335,7 +335,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB)) assert(rightKeys === Seq(exprC, exprD)) @@ -353,7 +353,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB)) assert(rightKeys === Seq(exprC, exprD)) assert(p.expressions == Seq(exprC)) @@ -372,7 +372,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB, exprB)) assert(rightKeys === Seq(exprA, exprC, exprC)) assert(p.expressions == Seq(exprA, exprC, exprA)) @@ -388,7 +388,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB, exprB)) assert(rightKeys === Seq(exprA, exprC, exprD)) assert(p.expressions == Seq(exprA, exprC, exprA)) @@ -444,8 +444,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.numPartitions == 5) assert(right.numPartitions == 5) case other => fail(other.toString) @@ -462,7 +462,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, DummySparkPlan(_, _, left: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.numPartitions == 10) assert(right.numPartitions == 10) assert(right.expressions == Seq(exprC, exprD)) @@ -481,8 +481,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.numPartitions == 5) assert(left.expressions == Seq(exprA, exprB)) assert(right.numPartitions == 5) @@ -492,7 +492,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, DummySparkPlan(_, _, left: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.numPartitions == 1) assert(right.numPartitions == 1) assert(right.expressions == Seq(exprC)) @@ -510,8 +510,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.numPartitions == conf.numShufflePartitions) assert(left.expressions == Seq(exprA, exprB)) assert(right.numPartitions == conf.numShufflePartitions) @@ -529,7 +529,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, DummySparkPlan(_, _, left: PartitioningCollection, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.numPartitions == 10) assert(right.numPartitions == 10) assert(right.expressions == Seq(exprA)) @@ -545,7 +545,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { Inner, None, plan1, plan2) applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, right: PartitioningCollection, _, _), _), _) => assert(left.numPartitions == 20) assert(left.expressions == Seq(exprC)) @@ -587,7 +587,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, DummySparkPlan(_, _, left: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB)) assert(right.expressions === Seq(exprC, exprD)) assert(left.numPartitions == 6) @@ -606,7 +606,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, right: HashPartitioning, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB)) assert(right.expressions === Seq(exprC, exprD)) @@ -624,7 +624,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, right: HashPartitioning, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB)) assert(right.expressions === Seq(exprC, exprD)) @@ -644,7 +644,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, right: HashPartitioning, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB)) assert(right.expressions === Seq(exprC, exprD)) @@ -666,8 +666,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB)) assert(right.expressions === Seq(exprC, exprD)) assert(left.numPartitions == conf.numShufflePartitions) @@ -691,7 +691,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, left: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB)) assert(rightKeys === Seq(exprC, exprD)) assert(left.numPartitions == 9) @@ -714,8 +714,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { var smjExec = SortMergeJoinExec(exprA :: Nil, exprC :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA)) assert(rightKeys === Seq(exprC)) assert(left.numPartitions == 20) @@ -733,7 +733,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA)) assert(rightKeys === Seq(exprC)) assert(right.numPartitions == 10) @@ -765,8 +765,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { } else { EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA)) assert(rightKeys === Seq(exprC)) assert(left.numPartitions == 5) @@ -875,8 +875,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB, exprB)) assert(right.expressions === Seq(exprA, exprC, exprC)) case other => fail(other.toString) @@ -951,8 +951,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: exprC :: Nil, exprA :: exprB :: exprC :: Nil, Inner, None, plan1, plan2) applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB, exprC)) assert(right.expressions === Seq(exprA, exprB, exprC)) case other => fail(other.toString) @@ -969,8 +969,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB, exprB)) assert(right.expressions === Seq(exprA, exprC, exprC)) case other => fail(other.toString) @@ -987,8 +987,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB, exprB)) assert(right.expressions === Seq(exprA, exprC, exprC)) case other => fail(other.toString) @@ -1008,8 +1008,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB, exprB)) assert(right.expressions === Seq(exprA, exprC, exprC)) case other => fail(other.toString) @@ -1130,7 +1130,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { case ShuffledHashJoinExec(_, _, _, _, _, DummySparkPlan(_, _, left: KeyedPartitioning, _, _), ShuffleExchangeExec(KeyedPartitioning(attrs, pks, _, _), - DummySparkPlan(_, _, SinglePartition, _, _), _, _), _) => + DummySparkPlan(_, _, SinglePartition, _, _), _, _, _), _) => assert(left.expressions == a1 :: Nil) assert(attrs == a1 :: Nil) assert(partitionKeys == pks.map(_.row)) @@ -1234,8 +1234,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(join) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(p1: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(p2: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(p1: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(p2: HashPartitioning, _, _, _, _), _), _) => // Both sides should be shuffled to default partitions assert(p1.numPartitions == 10) assert(p2.numPartitions == 10) @@ -1259,8 +1259,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(join) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), _) => // Both sides shuffled due to key mismatch case other => fail(s"Expected shuffles on both sides, but got: $other") } @@ -1278,7 +1278,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(join) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), SortExec(_, _, _: DummySparkPlan, _), _) => // Left side shuffled, right side kept as-is case other => fail(s"Expected shuffle on the left side, but got: $other") @@ -1296,8 +1296,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(join) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), _) => // Both sides shuffled due to canCreatePartitioning = false case other => fail(s"Expected shuffles on both sides, but got: $other") } @@ -1371,8 +1371,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(join) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(p1: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(p2: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(p1: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(p2: HashPartitioning, _, _, _, _), _), _) => // Both sides should be shuffled because partition keys not in join keys assert(p1.numPartitions == 10) assert(p2.numPartitions == 10) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala new file mode 100644 index 0000000000000..cf6abeff4884e --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala @@ -0,0 +1,223 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.streaming + +import java.io.{File, IOException} +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.atomic.AtomicInteger + +import org.apache.spark.SparkException +import org.apache.spark.scheduler.{SparkListener, SparkListenerStageCompleted} +import org.apache.spark.scheduler.SparkListenerStageSubmitted +import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec +import org.apache.spark.sql.execution.streaming.sources.{ContinuousMemorySink, LowLatencyMemoryStream} +import org.apache.spark.sql.execution.streaming.state.{FailureInjectionCheckpointFileManager, FailureInjectionFileSystem, FailureInjectionState} +import org.apache.spark.sql.functions.udf +import org.apache.spark.sql.internal.SQLConf + +/** Driver-side switch a UDF reads on executors to fail a task on demand (for fault-tolerance). */ +object RealTimeModePipelinedShuffleSuite { + @volatile var failTasks: Boolean = false +} + +/** + * Tests that a stateful Real-Time Mode query (streaming `dropDuplicates`) runs when its repartition + * shuffle is a PipelinedShuffleDependency, so the source-scan producer stage and the dedup consumer + * stage are co-scheduled and stream records through a transient shuffle instead of the consumer + * waiting for the producer to fully materialize. + * + * The Real-Time Mode operator allowlist check is left enabled, so the tests also confirm that the + * dedup plan's operators (StreamingDeduplicateExec, StateStoreRestore/Save, ShuffleExchangeExec) + * are admitted. They drive the source with the standard `testStream` DSL, which advances the query + * across multiple batches. + */ +class RealTimeModePipelinedShuffleSuite extends StreamRealTimeModeManualClockSuiteBase { + import testImplicits._ + + override def beforeEach(): Unit = { + super.beforeEach() + RealTimeModePipelinedShuffleSuite.failTasks = false + } + + test("stateful dedup runs in Real-Time Mode over a pipelined shuffle") { + // Track, from the driver, whether the producer (source scan) and consumer (dedup) stages of the + // pipelined group were ever RUNNING simultaneously. runningStages holds currently-running stage + // ids (added on submit, removed on completion); maxConcurrentStages records the peak. A + // sequential producer-then-consumer schedule never exceeds one running stage at a time. + val runningStages = ConcurrentHashMap.newKeySet[Int]() + val maxConcurrentStages = new AtomicInteger(0) + val listener = new SparkListener { + override def onStageSubmitted(e: SparkListenerStageSubmitted): Unit = { + runningStages.add(e.stageInfo.stageId) + maxConcurrentStages.accumulateAndGet(runningStages.size(), Math.max) + } + override def onStageCompleted(e: SparkListenerStageCompleted): Unit = { + runningStages.remove(e.stageInfo.stageId) + } + } + spark.sparkContext.addSparkListener(listener) + + try { + val inputData = LowLatencyMemoryStream[(String, Int)] + + // scan --shuffle(repartition by key)--> streaming dropDuplicates --> sink. + // dropDuplicates on "key" forces a hash-partitioning ShuffleExchangeExec, which the + // IncrementalExecution rule marks pipelined for a Real-Time Mode batch. + val result = inputData + .toDF() + .select($"_1".as("key")) + .dropDuplicates("key") + .select($"key") + + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + // Batch 0: three distinct keys, each sent twice -> dedup emits each once. A Real-Time Mode + // batch runs for a fixed duration and emits in real time, so CheckAnswerWithTimeout polls + // the sink (rather than blocking for batch completion) and advanceRealTimeClock ends it. + AddData(inputData, ("a", 1), ("b", 1), ("c", 1), ("a", 2), ("b", 2), ("c", 2)), + StartStream(), + CheckAnswerWithTimeout(60000, "a", "b", "c"), + advanceRealTimeClock, + WaitUntilBatchProcessed(0), + // Batch 1: all duplicates of already-seen keys plus one new key -> only "d" is new. + AddData(inputData, ("a", 3), ("b", 3), ("c", 3), ("d", 1)), + CheckAnswerWithTimeout(60000, "a", "b", "c", "d"), + // Every Real-Time Mode shuffle exchange in the executed plan is pipelined, and the producer + // + consumer stages were genuinely co-scheduled (>= 2 stages ran at once). + Execute { q => + val executedPlan = q.lastExecution.executedPlan + val exchanges = executedPlan.collect { case s: ShuffleExchangeExec => s } + assert(exchanges.nonEmpty, "expected a shuffle exchange in the dedup plan") + assert(exchanges.forall(_.pipelined), + "expected all Real-Time Mode shuffle exchanges to be pipelined, got: " + + exchanges.map(e => s"pipelined=${e.pipelined}").mkString(", ")) + assert(maxConcurrentStages.get() >= 2, + s"expected >= 2 stages running concurrently, saw max ${maxConcurrentStages.get()}") + }, + StopStream + ) + } finally { + spark.sparkContext.removeSparkListener(listener) + } + } + + test("dedup over a pipelined shuffle recovers from a task failure via checkpoint restart") { + withTempDir { checkpointDir => + // A UDF that throws on demand, to fail a task mid-batch. Placed after the dedup so the + // failure lands in the pipelined consumer stage while the query is running. + val failUDF = udf { (key: String) => + if (RealTimeModePipelinedShuffleSuite.failTasks) { + throw new RuntimeException(s"forced task failure on $key") + } + key + } + + val inputData = LowLatencyMemoryStream[(String, Int)] + val result = inputData + .toDF() + .select($"_1".as("key")) + .dropDuplicates("key") + .select(failUDF($"key").as("key")) + + // First run: dedup "a","b" (batch 0 commits), then a task failure fails the query. RTM does + // not retry tasks, so a single task failure fails the whole batch/query. + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, ("a", 1), ("b", 1), ("a", 2)), + StartStream(checkpointLocation = checkpointDir.getAbsolutePath), + CheckAnswerWithTimeout(60000, "a", "b"), + advanceRealTimeClock, + WaitUntilBatchProcessed(0), + Execute { _ => RealTimeModePipelinedShuffleSuite.failTasks = true }, + AddData(inputData, ("c", 1)), + advanceRealTimeClock, + ExpectFailure[SparkException] { ex => + val msg = Option(ex.getCause).map(_.getMessage).getOrElse(ex.getMessage) + assert(msg != null && msg.contains("forced task failure"), + s"expected a forced task failure, got: $msg") + } + ) + + // Restart from the same checkpoint. The dedup state from batch 0 must survive: "a" and "b" + // are already seen and must NOT be re-emitted; only the genuinely-new "c","d" appear. This + // proves recovery-to-last-committed-batch works when the shuffle is pipelined. + RealTimeModePipelinedShuffleSuite.failTasks = false + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + StartStream(checkpointLocation = checkpointDir.getAbsolutePath), + AddData(inputData, ("a", 3), ("b", 3), ("c", 3), ("d", 1)), + CheckAnswerWithTimeout(60000, "c", "d"), + advanceRealTimeClock, + StopStream + ) + } + } + + /** Run `f` with a temp dir whose checkpoint file ops can be fault-injected via injectionState. */ + private def withTempDirAllowFailureInjection(f: (File, FailureInjectionState) => Unit): Unit = { + withTempDir { dir => + val injectionState = FailureInjectionFileSystem.registerTempPath(dir.getPath) + try { + f(dir, injectionState) + } finally { + FailureInjectionFileSystem.removePathFromTempToInjectionState(dir.getPath) + } + } + } + + test("dedup over a pipelined shuffle recovers from a commit-log write failure") { + withSQLConf( + SQLConf.STREAMING_CHECKPOINT_FILE_MANAGER_CLASS.parent.key -> + classOf[FailureInjectionCheckpointFileManager].getName) { + withTempDirAllowFailureInjection { (checkpointDir, injectionState) => + val inputData = LowLatencyMemoryStream[(String, Int)] + val result = inputData + .toDF() + .select($"_1".as("key")) + .dropDuplicates("key") + .select($"key") + + // Batch 0 dedups "a","b" and commits. Then fail the close() of batch 1's commit-log write, + // so batch 1 cannot commit and the query fails after processing. + injectionState.createAtomicDelayCloseRegex = Seq(".*/commits/1") + + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, ("a", 1), ("b", 1)), + StartStream(checkpointLocation = checkpointDir.getAbsolutePath), + CheckAnswerWithTimeout(60000, "a", "b"), + advanceRealTimeClock, + WaitUntilBatchProcessed(0), + AddData(inputData, ("c", 1)), + CheckAnswerWithTimeout(60000, "a", "b", "c"), + advanceRealTimeClock, + // The injected close() failure surfaces as the underlying IOException failing the batch. + ExpectFailure[IOException]() + ) + + // Clear the injection and restart from the same checkpoint. Batch 0's dedup state must + // survive (a, b already seen); the uncommitted batch 1 is re-run, so its new key "c" is + // still emitted, and a further new key "d" is added. + injectionState.createAtomicDelayCloseRegex = Seq.empty + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + StartStream(checkpointLocation = checkpointDir.getAbsolutePath), + AddData(inputData, ("a", 2), ("b", 2), ("c", 2), ("d", 1)), + CheckAnswerWithTimeout(60000, "c", "d"), + advanceRealTimeClock, + StopStream + ) + } + } + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala index c46f0076721b9..c0983f338abe5 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala @@ -677,8 +677,8 @@ abstract class StreamingInnerJoinBase extends StreamingJoinSuite { assert(query.lastExecution.executedPlan.collect { case j @ StreamingSymmetricHashJoinExec(_, _, _, _, _, _, _, _, _, - ShuffleExchangeExec(opA: HashPartitioning, _, _, _), - ShuffleExchangeExec(opB: HashPartitioning, _, _, _), _) + ShuffleExchangeExec(opA: HashPartitioning, _, _, _, _), + ShuffleExchangeExec(opB: HashPartitioning, _, _, _, _), _) if partitionExpressionsColumns(opA.expressions) === Seq("a", "b") && partitionExpressionsColumns(opB.expressions) === Seq("a", "b") && opA.numPartitions == numPartitions && opB.numPartitions == numPartitions => j