diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index 875be3d75158..4104d04df766 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -1126,6 +1126,12 @@ ], "sqlState" : "0A000" }, + "CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT" : { + "message" : [ + "Cannot run the pipelined stage group: it needs concurrent task slots to run all its stages together, but only are currently free. The stages of a pipelined group must run concurrently, so the cluster must have enough free slots for all of them at once. Provision more executors, or reduce the parallelism of the query." + ], + "sqlState" : "53000" + }, "CONCURRENT_STREAM_LOG_UPDATE" : { "message" : [ "Concurrent update to the log. Multiple streaming jobs detected for .", @@ -6313,6 +6319,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/internal/config/package.scala b/core/src/main/scala/org/apache/spark/internal/config/package.scala index c8bd6410be27..2f20efbc0a0a 100644 --- a/core/src/main/scala/org/apache/spark/internal/config/package.scala +++ b/core/src/main/scala/org/apache/spark/internal/config/package.scala @@ -2006,6 +2006,21 @@ package object config { .checkValue(v => v > 0, "The max failures should be a positive value.") .createWithDefault(40) + private[spark] val PIPELINED_GROUP_SLOT_CHECK_ENABLED = + ConfigBuilder("spark.scheduler.pipelinedGroup.slotCheck.enabled") + .internal() + .doc("When true, before co-scheduling a pipelined-shuffle stage group the DAGScheduler " + + "checks that the group's total task demand fits in the currently free slots of its " + + "resource profile (total capacity minus the outstanding -- running plus enqueued -- " + + "tasks of other work), and fails the job with CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT " + + "rather than co-scheduling a group that cannot fit and deadlocking. Set to false for " + + "deployments that admit capacity out-of-band (e.g. a slot reservation), which then own " + + "admission. Only applies to jobs that use a pipelined shuffle dependency.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .booleanConf + .createWithDefault(true) + private[spark] val NUM_CANCELLED_JOB_GROUPS_TO_TRACK = ConfigBuilder("spark.scheduler.numCancelledJobGroupsToTrack") .doc("The maximum number of tracked job groups that are cancelled with " + 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 9876668194a8..4491461ee119 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 5e608308de29..ee820b1b9663 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.{ResourceProfile, TaskResourceProfile} import org.apache.spark.resource.ResourceProfile.{DEFAULT_RESOURCE_PROFILE_ID, EXECUTOR_CORES_LOCAL_PROPERTY, PYSPARK_MEMORY_LOCAL_PROPERTY} import org.apache.spark.rpc.RpcTimeout @@ -169,6 +169,29 @@ private[spark] class DAGScheduler( // Stages that must be resubmitted due to fetch failures private[scheduler] val failedStages = new HashSet[Stage] + /** + * Deferred completion for pipelined groups. When a stage is co-scheduled with a pipelined + * producer that is still running (a "pipelined consumer"), the stage/job-completion decision from + * its successful task-completion events must not be processed yet: doing so could finish the + * consumer's job and cancel the still-running producer, or make the consumer's output observable + * before the producer's. We buffer such events here and replay them once every pipelined producer + * the consumer was waiting on has finished. + * + * Keyed by the consumer stage. `parents` is the set of its pipelined producer stages not yet + * finished; `delayedTaskCompletionEvents` are its Success CompletionEvents held until then. + * Entries exist only for stages that were co-scheduled with a not-yet-finished pipelined + * producer, so this is empty for any job without a pipelined dependency. + */ + private[scheduler] case class DependentStageInfo( + parents: HashSet[Stage] = new HashSet[Stage], + 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 @@ -604,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 => @@ -653,6 +691,7 @@ private[spark] class DAGScheduler( checkBarrierStageWithDynamicAllocation(rdd) checkBarrierStageWithNumSlots(rdd, resourceProfile) checkBarrierStageWithRDDChainPattern(rdd, rdd.getNumPartitions) + checkPipelinedProducerSupported(shuffleDep) val numTasks = rdd.partitions.length val parents = getOrCreateParentStages(shuffleDeps, jobId) val id = nextStageId.getAndIncrement() @@ -676,6 +715,63 @@ private[spark] class DAGScheduler( 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 @@ -945,6 +1041,223 @@ private[spark] class DAGScheduler( missing.toList } + /** + * Whether the RDD graph rooted at `finalRDD` contains a [[PipelinedShuffleDependency]] anywhere. + * Walks the RDD dependency graph directly (not the stage graph), so it can be checked before any + * stages are created -- letting a job be rejected up front without leaving partial scheduler + * state behind. + */ + private def rddGraphHasPipelinedDependency(finalRDD: RDD[_]): Boolean = { + // traverseRDDGraphUntil stops and returns false as soon as the visitor returns false; we use + // that to short-circuit on the first pipelined dependency found. It returns true if the whole + // graph was visited without stopping (i.e. none found), so negate the result. + !traverseRDDGraphUntil(finalRDD) { (rdd, enqueue) => + val hasPipelined = rdd.dependencies.exists { + case _: PipelinedShuffleDependency[_, _, _] => true + case dep => + enqueue(dep.rdd) + false + } + !hasPipelined // keep traversing while none found; stop (return false) when one is found + } + } + + /** + * Classifies the shuffle boundaries in the RDD graph rooted at `finalRDD` by kind, walking the + * RDD dependency graph directly (before any stages are created). Returns whether the graph + * contains any [[PipelinedShuffleDependency]] and whether it contains any regular (non-pipelined) + * `ShuffleDependency`. Narrow dependencies are not boundaries and are ignored. + * + * A job must be either ALL-regular or ALL-pipelined: a job whose shuffle graph mixes a pipelined + * shuffle with a regular one is rejected fail-fast (see `handleJobSubmitted`). Restricting to a + * single kind makes the whole job one pipelined group when pipelined, so gang admission can be + * decided up front before any stage is submitted. + */ + private def classifyJobShuffleKinds(finalRDD: RDD[_]): (Boolean, Boolean) = { + var hasPipelined = false + var hasRegular = false + traverseRDDGraph(finalRDD) { (rdd, enqueue) => + rdd.dependencies.foreach { dep => + dep match { + case _: PipelinedShuffleDependency[_, _, _] => hasPipelined = true + case _: ShuffleDependency[_, _, _] => hasRegular = true + case _ => // narrow dependency: not a boundary + } + // Descend through every edge (shuffle and narrow) so a pipelined boundary behind a regular + // one -- or vice versa -- anywhere in the graph is still detected. traverseRDDGraph dedups. + enqueue(dep.rdd) + } + } + (hasPipelined, hasRegular) + } + + /** + * Reject a job that uses a pipelined shuffle in combination with a cluster feature that a + * pipelined group cannot support. Checked up front, before any stage is created, so a rejection + * leaves no partial scheduler state. Used by the result-job path (handleJobSubmitted); the + * map-stage-job path rejects a pipelined dependency outright (see handleMapStageSubmitted), which + * subsumes these. Returns true (and fails the job via `listener`) if rejected; false otherwise. + * The RDD-graph walk runs only when a relevant feature is enabled and is inert for jobs without a + * pipelined dependency. + * + * Rejected combinations: + * - Speculation: a speculative copy of a producer would race a consumer already reading the + * producer's partial output, with no commit barrier protecting the read. + * - Dynamic allocation: a pipelined group is gang-scheduled (all member stages must run at + * once), and the free-slot admission check measures currently-active executors. Under dynamic + * allocation a job can start before any executor has spun up, so the group would be failed + * with CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT against a transient 0-slot snapshot even though + * the cluster would soon have capacity. Barrier scheduling forbids the same combination + * (checkBarrierStageWithDynamicAllocation); pipelined groups do likewise. + */ + private def rejectUnsupportedPipelinedJob( + jobId: Int, + finalRDD: RDD[_], + listener: JobListener): Boolean = { + // Only walk the RDD graph if a relevant feature is on, then only once. + val speculationOn = sc.conf.get(config.SPECULATION_ENABLED) + val dynAllocOn = Utils.isDynamicAllocationEnabled(sc.conf) + if ((speculationOn || dynAllocOn) && rddGraphHasPipelinedDependency(finalRDD)) { + if (speculationOn) { + logWarning(log"Rejecting job ${MDC(JOB_ID, jobId)}: speculation is incompatible with " + + log"pipelined shuffle dependencies") + listener.jobFailed(new SparkException( + "Speculative execution is not supported for a job that uses a pipelined shuffle " + + s"dependency. Disable ${config.SPECULATION_ENABLED.key} for such jobs.")) + return true + } + // dynAllocOn + logWarning(log"Rejecting job ${MDC(JOB_ID, jobId)}: dynamic allocation is incompatible " + + log"with pipelined shuffle dependencies") + listener.jobFailed(new SparkException( + "Dynamic allocation is not supported for a job that uses a pipelined shuffle dependency: " + + "a pipelined stage group must be co-scheduled against a known cluster size. Disable " + + s"${config.DYN_ALLOCATION_ENABLED.key} for such jobs.")) + return true + } + 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 @@ -1016,6 +1329,23 @@ private[spark] class DAGScheduler( logDebug("Removing stage %d from failed set.".format(stageId)) failedStages -= stage } + // Drop any pipelined-completion deferral keyed on this stage (as consumer), and + // remove it from other consumers' pending-producer sets (as producer), so no + // deferral outlives its job (e.g. on job abort before the producers finished). + // A consumer's buffered completion events hold TaskEnds that have not been emitted + // yet (the whole event is deferred until group completion). This teardown is a job + // failure / cancellation, so the buffered successes must NOT be applied as results + // -- but their TaskEnds must still be flushed, exactly as the producer-failed drop + // path does (see releaseDeferredPipelinedConsumers), or a listener tracking active + // tasks (e.g. AppStatusListener) would leak these tasks as perpetually running. + // Normally releaseDeferredPipelinedConsumers has already drained the deferral by + // the time cleanup runs (cancelRunningIndependentStages finishes each running + // producer first); this covers the case where it has not -- e.g. a producer left in + // resubmit limbo (finished but not available) that cancelRunningIndependentStages + // skips because it is neither running nor failed. + dependentStageMap.remove(stage).foreach(_.delayedTaskCompletionEvents.foreach( + postTaskEnd)) + dependentStageMap.values.foreach(_.parents -= stage) } // data structures based on StageId stageIdToStage -= stageId @@ -1346,6 +1676,172 @@ private[spark] class DAGScheduler( } } + /** + * Reconsider waiting stages that read `runningParent` through a pipelined edge, now that it has + * started running. A pipelined edge is non-sequencing, so such a consumer can be co-scheduled as + * soon as its producer is running -- it need not wait for the producer to finish. This is the + * "producer started" analog of `submitWaitingChildStages` (which fires on producer *completion*): + * without it, a pipelined consumer parked because its producer was not yet runnable (e.g. the + * producer sat behind a regular shuffle) would not be co-scheduled until the producer finished, + * losing the pipelining. Inert unless `runningParent` is the pipelined producer of a waiting + * stage. + */ + private def submitWaitingPipelinedChildStages(runningParent: Stage): Unit = { + val pipelinedChildren = waitingStages.filter { child => + child.parents.contains(runningParent) && isPipelinedProducer(runningParent) + }.toArray + // Remove them from waitingStages before resubmitting, or submitStage's `!waitingStages(stage)` + // guard would treat them as already-scheduled and no-op (mirrors submitWaitingChildStages). + waitingStages --= pipelinedChildren + for (child <- pipelinedChildren.sortBy(_.firstJobId)) { + logInfo(log"Reconsidering ${MDC(STAGE, child)} now that its pipelined producer " + + log"${MDC(STAGE, runningParent)} is running") + submitStage(child) + } + } + + /** + * Whether `stage` produces its output through a [[PipelinedShuffleDependency]] -- i.e. it is a + * pipelined producer, whose consumers may run concurrently with it. Callers that need the + * producer/consumer relationship check `child.parents.contains(stage)` separately. + */ + private def isPipelinedProducer(stage: Stage): Boolean = stage match { + case m: ShuffleMapStage => m.shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]] + case _ => false + } + + /** + * The total concurrent-task demand of an all-pipelined job, computed from the RDD graph BEFORE + * any stage is created (so a rejection based on it leaves no partial scheduler state, exactly as + * the barrier slot check and the speculation/DA reject do). Because a job is either all-regular + * or all-pipelined, an all-pipelined job's whole stage graph is one pipelined group; its members + * are the final result stage plus every pipelined producer. Each member's task count is its RDD's + * partition count (`rdd.partitions.length`), matching how `createShuffleMapStage` derives + * `numTasks`. `finalNumPartitions` is the result stage's task count (the number of partitions the + * job runs, which may be a subset of `finalRDD.partitions`). + */ + private def pipelinedJobConcurrentTaskDemand(finalRDD: RDD[_], finalNumPartitions: Int): Int = { + var demand = finalNumPartitions + traverseRDDGraph(finalRDD) { (rdd, enqueue) => + rdd.dependencies.foreach { + case pd: PipelinedShuffleDependency[_, _, _] => demand += pd.rdd.partitions.length + case _ => // regular/narrow deps do not occur in an all-pipelined job's group + } + rdd.dependencies.foreach(dep => enqueue(dep.rdd)) + } + demand + } + + /** + * Up-front gang admission for an all-pipelined job, checked BEFORE any stage exists. Because a + * job is either all-regular or all-pipelined (mixed jobs are already rejected), an all-pipelined + * job's whole stage graph is one pipelined group with no regular prefix, so its full demand is + * known up front and the group is ready to admit immediately. Checking here (rather than in + * `submitStage` once a producer is already running) is true all-or-nothing gang admission: the + * whole group is admitted, or the job is failed before any member runs, so a member is never left + * running while a sibling cannot get slots -- and, like the barrier slot check, a rejection + * leaves no partial scheduler state. + * + * Free-slot accounting: demand vs. total capacity (`maxNumConcurrentTasks` for the default + * profile -- a group is required to be single-profile) minus what OTHER work (other jobs) has + * outstanding -- running plus enqueued -- in that profile. Counting enqueued, not just running, + * tasks charges a busy neighbor's queued backlog against capacity too, so a group is not admitted + * against slots other work is already committed to. Fails the job via `listener` with no retry -- + * a transient shortfall is the caller's to retry (e.g. the streaming batch loop reruns it). + * Returns true (job failed) if it does not fit. + * + * The likeness to barrier's slot check is only that both reject before any stage is created + * (leaving no partial state) and compute demand from the RDD graph. Retry behavior differs + * deliberately: barrier RE-POSTS the job and re-runs its check on a timer up to + * `spark.scheduler.barrier.maxConcurrentTasksCheck.maxFailures` times; pipelined admission is + * TERMINAL (one check, then fail), delegating transient-shortfall retry to the caller. + * + * The check can be turned off with `spark.scheduler.pipelinedGroup.slotCheck.enabled=false` (for + * deployments that admit capacity out-of-band, e.g. via a slot reservation). + */ + private def rejectUnadmittablePipelinedGroup( + jobId: Int, finalRDD: RDD[_], partitions: Array[Int], listener: JobListener): Boolean = { + 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 + val demand = pipelinedJobConcurrentTaskDemand(finalRDD, partitions.length) + val totalSlots = maxConcurrentTasksForProfile(rp.id) + // No stage of this job exists yet, so it has no outstanding tasks of its own to exclude; only + // other concurrent jobs' outstanding demand is charged, resource-profile-scoped. + val outstandingForOthers = outstandingTasksForOtherWork(rp.id, Set.empty) + val freeSlots = math.max(0, totalSlots - outstandingForOthers) + if (demand > freeSlots) { + logWarning(log"Rejecting job ${MDC(JOB_ID, jobId)}: pipelined stage group needs " + + log"${MDC(NUM_TASKS, demand)} concurrent task slots but only " + + log"${MDC(NUM_SLOTS, freeSlots)} are free") + listener.jobFailed(new SparkException( + errorClass = "CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT", + messageParameters = scala.collection.immutable.Map( + "numTasks" -> demand.toString, "numSlots" -> freeSlots.toString), + cause = null)) + true + } else { + false + } + } + + /** + * The cluster's total concurrent-task capacity for the given resource profile. Extracted as a + * seam so tests can control it without changing the cluster's core count. Production reads it + * from the scheduler backend, exactly as barrier's slot check does. + */ + protected def maxConcurrentTasksForProfile(rpId: Int): Int = { + val rp = sc.resourceProfileManager.resourceProfileFromId(rpId) + sc.maxNumConcurrentTasks(rp) + } + + /** + * Outstanding tasks (running plus enqueued) in the given resource profile, for work OTHER than + * the stages in `excludeStageIds`. Counting enqueued tasks, not just running ones, means a busy + * neighbor's queued backlog is charged against capacity too, so a group is not admitted against + * slots that other work is already committed to using. Resource-profile-scoped to match + * `maxConcurrentTasksForProfile`. Extracted as a seam so tests can control occupancy without + * launching real tasks; returns 0 for a non-TaskSchedulerImpl backend. + */ + protected def outstandingTasksForOtherWork(rpId: Int, excludeStageIds: Set[Int]): Int = + taskScheduler match { + case impl: TaskSchedulerImpl => + impl.outstandingTasksForOtherWorkInProfile(rpId, excludeStageIds) + case _ => 0 + } + + /** + * Whether `stage` is a member of a pipelined group -- i.e. it is connected to another stage by a + * [[PipelinedShuffleDependency]], either as the producer (it writes such a shuffle) or as a + * consumer (one of its direct parent shuffle dependencies is pipelined). Members run + * concurrently over a transient shuffle that cannot be re-read in isolation, so a member's task + * failure must fail the whole group rather than resubmit one stage; the task scheduler keys its + * fail-fast behavior off this (via TaskSet.isPipelined). False for any stage in a job with no + * pipelined dependency. + * + * NOTE: this is deliberately defined here alongside the other group-topology helper + * (`isPipelinedProducer`). Its call sites are the group-atomic failure handling: tagging the + * member's `TaskSet.isPipelined` at submission and routing a member FetchFailed to a whole-group + * abort. + */ + 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 // stage the one with the highest priority (highest-priority pool, earliest created). @@ -1489,8 +1985,52 @@ private[spark] class DAGScheduler( return } + // A job that uses a pipelined shuffle co-schedules its producer and consumer stages, which is + // incompatible with speculation (a speculative producer copy would race a consumer reading its + // partial output) and with dynamic allocation (the gang-scheduling slot check would fail + // against a not-yet-warmed cluster). Reject such a job up front -- before any stages are + // created, so no partial scheduler state is left behind. (Inert for jobs without any pipelined + // dependency.) + if (rejectUnsupportedPipelinedJob(jobId, finalRDD, listener)) { + return + } + + // A job must be either ALL-regular or ALL-pipelined, not a mix: a job whose shuffle graph + // combines a pipelined shuffle with a regular one is rejected up front (before any stage is + // created). Restricting to a single kind makes an all-pipelined job's whole stage graph one + // pipelined group with no regular prefix, so gang admission is decided up front (see the slot + // check below) rather than mid-DAG once a producer is already running. + val (hasPipelined, hasRegular) = classifyJobShuffleKinds(finalRDD) + if (hasPipelined && hasRegular) { + logWarning(log"Rejecting job ${MDC(JOB_ID, jobId)}: a job mixing a pipelined shuffle with " + + log"a regular shuffle is not supported") + listener.jobFailed(new SparkException( + "A job that mixes a pipelined shuffle dependency with a regular shuffle dependency is " + + "not supported: a job must be either all-regular or all-pipelined.")) + return + } + + // Gang admission for an all-pipelined job: the whole stage graph is one pipelined group, so + // check up front (before any stage is created) that the cluster can run the entire group + // concurrently. If it cannot fit, fail the job now -- no partial scheduler state, and no member + // ever left running while a sibling waits on slots (true all-or-nothing gang admission). Inert + // for a regular job (no pipelined dependency). + 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) @@ -1523,6 +2063,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) @@ -1532,6 +2081,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)}) " + @@ -1564,6 +2121,24 @@ private[spark] class DAGScheduler( listener: JobListener, artifacts: JobArtifactSet, properties: Properties): Unit = { + // A map-stage job (SparkContext.submitMapStage, e.g. AQE's ShuffleExchangeExec) materializes a + // shuffle to produce its map-output statistics. A PipelinedShuffleDependency cannot serve that: + // it is transient with no durable, addressable map output, so it registers no map outputs and + // getStatistics would return all-zero stats (misleading the map-stage/AQE consumer); and its + // producer/consumer co-scheduling makes speculation unsafe. Reject a pipelined + // dependency submitted as a map-stage job outright, up front, before any stage is created so no + // partial scheduler state is left behind. Inert for a regular ShuffleDependency. (The + // result-job path, handleJobSubmitted, rejects the speculation and dynamic-allocation cases via + // rejectUnsupportedPipelinedJob, but a pipelined dependency is otherwise a legitimate internal + // edge there.) + if (dependency.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) { + logWarning(log"Rejecting map-stage job ${MDC(JOB_ID, jobId)}: a pipelined shuffle dependency " + + log"cannot be materialized as a map-stage job") + listener.jobFailed(new SparkException( + "A pipelined shuffle dependency cannot be submitted as a map-stage job: it has no durable " + + "map output to produce statistics from. This is not supported.")) + return + } // Submitting this map stage might still require the creation of some parent stages, so make // sure that happens. var finalStage: ShuffleMapStage = null @@ -1630,11 +2205,81 @@ private[spark] class DAGScheduler( logInfo(log"Submitting ${MDC(STAGE, stage)} (${MDC(RDD_ID, stage.rdd)}), " + log"which has no missing parents") submitMissingTasks(stage, jobId.get) + // If this stage is the pipelined producer of a waiting consumer, co-schedule it now. + submitWaitingPipelinedChildStages(stage) } else { for (parent <- missing) { submitStage(parent) } - waitingStages += stage + + // Submitting a parent can abort the job during this recursion (e.g. a parent that has + // exhausted its stage attempts hits abortStage, which cleans up all of the job's stages + // including this one and fails the job). If that happened, `stage` is no longer + // registered; do not proceed to co-schedule or re-park it (re-parking would re-insert a + // job-less stage into waitingStages -- a scheduler-state leak). Inert for a + // non-aborting submit. (Pipelined group admission is decided up front in + // handleJobSubmitted, so it cannot abort the group from within this recursion.) + if (!stageIdToStage.contains(stage.id)) { + logInfo(log"${MDC(STAGE, stage)} was removed during parent submission (its job was " + + log"aborted); not co-scheduling or re-parking it") + return + } + + // A missing parent reached through a PipelinedShuffleDependency ("pipelined parent") + // is incrementally readable: this stage may run before that parent materializes, so + // the two are co-scheduled. `missing` already holds the direct parent shuffle-map + // stages (from getMissingParentStages), so classify them by their shuffle dependency + // type -- no extra graph walk. For a job with no pipelined dependency, pipelinedMissing + // is empty and this stage simply parks in waitingStages, exactly as before. + val (pipelinedMissing, regularMissing) = missing.partition(isPipelinedProducer) + // Co-schedule only if EVERY missing parent is pipelined AND each is actually running + // now. submitStage above may have parked a pipelined parent in waitingStages (e.g. it + // has its own regular missing parent); running this stage against a not-yet-running + // producer would strand it. If any parent is regular or not yet runnable, park this + // stage. It is then resubmitted and co-scheduled once its parents become runnable -- + // via submitWaitingChildStages when a regular parent completes, or via + // submitWaitingPipelinedChildStages when a pipelined parent starts running. + 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). 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, + // so its successful completions are deferred until those producers finish. + val deferral = + dependentStageMap.getOrElseUpdate(stage, DependentStageInfo()) + deferral.parents ++= pipelinedMissing + submitMissingTasks(stage, jobId.get) + // This stage is now running; if it is itself the pipelined producer of a waiting + // consumer, co-schedule that consumer too. + submitWaitingPipelinedChildStages(stage) + } else { + waitingStages += stage + } } } } @@ -2018,9 +2663,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 @@ -2425,6 +3074,34 @@ private[spark] class DAGScheduler( val stage = stageIdToStage(task.stageId) + // Group-observable completion (spec S5): if this stage is a pipelined consumer co-scheduled + // with a still-running pipelined producer, defer its *successful* completion in full until the + // producer(s) finish. Buffer the whole CompletionEvent and return before ANY of its side + // effects run. Buffering the whole event DEFERS all three of those effects -- the accumulator + // update, the task-end listener event, and stage/job completion -- until the event is replayed; + // it is NOT just the stage/job-completion bookkeeping that waits. (A listener-visible + // consequence: a consumer task that finishes ahead of its producer emits no SparkListenerTaskEnd + // until the producer finishes, so the Spark UI reports that already-succeeded task as still + // running for the producer's remaining lifetime.) Else a consumer finishing ahead of its + // producer would advance job completion and cancel the still-running producer (via + // cancelRunningIndependentStages), or expose its output early. + // Deferring the entire event is what makes the side effects run exactly ONCE, at replay: + // releaseDeferredPipelinedConsumers re-posts the buffered event once the last producer completes + // (or drops it if a producer fails, S6), and it then re-enters here and runs the side effects + // normally. This deferral check must therefore precede updateAccumulators and postTaskEnd. Inert + // for jobs with no pipelined dependency (the map is empty), so the regular path is unchanged. + if (event.reason == Success) { + dependentStageMap.get(stage) match { + case Some(deferral) if deferral.parents.nonEmpty => + logInfo(log"Deferring completion of task ${MDC(TASK_ID, event.taskInfo.taskId)} in " + + log"pipelined consumer ${MDC(STAGE, stage)} until its producer(s) " + + log"${MDC(MISSING_PARENT_STAGES, deferral.parents.toSeq)} finish") + deferral.delayedTaskCompletionEvents += event + return + case _ => + } + } + // Make sure the task's accumulators are updated before any other processing happens, so that // we can post a task end event before any jobs or stages are updated. The accumulators are // only updated in certain cases. @@ -2533,7 +3210,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 @@ -2601,6 +3303,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) @@ -2717,39 +3446,7 @@ private[spark] class DAGScheduler( } // 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 => @@ -3289,6 +3986,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. @@ -3450,8 +4193,12 @@ private[spark] class DAGScheduler( /** * Marks a stage as finished and removes it from the list of running stages. + * + * `private[scheduler]` (not `private`) so tests can drive the pipelined-consumer release/retain + * decision below directly, including the retained branch (a not-yet-available pipelined producer + * about to resubmit), which is otherwise brittle to reach through the mock backend. */ - private def markStageAsFinished( + private[scheduler] def markStageAsFinished( stage: Stage, errorMessage: Option[String] = None, willRetry: Boolean = false): Unit = { @@ -3480,6 +4227,83 @@ private[spark] class DAGScheduler( } listenerBus.post(SparkListenerStageCompleted(stage.latestInfo)) runningStages -= stage + + // Release any pipelined consumers whose completion was deferred while this stage (a pipelined + // producer) was running. Only act when the producer's outcome is final: + // - willRetry: the stage is being retried, not finished -- leave consumers deferred. + // - a ShuffleMapStage that finished "successfully" (no errorMessage) but is NOT yet available + // (e.g. a bogus-epoch task left an output missing) is about to be resubmitted by + // processShuffleMapStageCompletion -- it is not truly done, so do NOT replay its consumers + // against soon-to-be-recomputed output; the release happens when its reattempt completes and + // it becomes available. + // Otherwise: producerFailed = the stage failed (errorMessage set) -> drop the consumers' + // buffered successes; producer succeeded -> replay them. + if (!willRetry) { + val producerFailed = errorMessage.isDefined + val producerAboutToResubmit = stage match { + case m: ShuffleMapStage => !producerFailed && !m.isAvailable + case _ => false + } + if (!producerAboutToResubmit) { + releaseDeferredPipelinedConsumers(stage, producerFailed = producerFailed) + } + } + } + + /** + * Called when a stage finishes. If `finishedStage` is a pipelined producer that some co-scheduled + * consumer was deferred on, resolve that consumer's deferral: + * - Producer FAILED: drop the consumer's buffered completions immediately, regardless of whether + * it still has other pending producers. A pipelined group is failed as a unit (S6), so once a + * producer fails the consumer's buffered successes must not be applied -- and dropping now, + * rather than remembering the failure until the last producer finishes, means no failure state + * outlives this call (a surviving producer that later finishes just finds the consumer gone). + * - Producer SUCCEEDED: remove it from the consumer's pending-producer set, and replay the + * buffered completions only once the LAST producer has succeeded (the set becomes empty). + * Inert unless `finishedStage` is a tracked producer. Recovery from a drop, if any, is a caller + * rerun of the whole group (a new job, S6) -- the scheduler does not rerun the group in place. + */ + private def releaseDeferredPipelinedConsumers( + finishedStage: Stage, producerFailed: Boolean): Unit = { + if (dependentStageMap.isEmpty) { + return + } + // Consumers waiting on this producer. + val released = dependentStageMap.filter { + case (_, d) => d.parents.contains(finishedStage) + }.keys.toArray + for (consumer <- released) { + val deferral = dependentStageMap(consumer) + deferral.parents -= finishedStage + if (producerFailed) { + // Drop immediately: the group is being failed as a unit, so this consumer's buffered + // successes must not be applied no matter what its other producers do. Removing the entry + // now leaves no deferral state to go stale for a later re-co-scheduling of this stage. + // This treats producerFailed as terminal, which holds for every producer-failure path that + // reaches here EXCEPT a barrier stage that fails a task and is resubmitted (its handler + // calls markStageAsFinished(errorMessage) without willRetry=true, then resubmits): + // there, dropping a fan-in consumer's deferral for a producer that will retry loses that + // consumer's buffered results. A barrier stage in a pipelined group is rejected up front by + // a later PR in this stack (checkPipelinedProducerSupported), which removes this path; on + // this PR standalone it is a known, narrow limitation (barrier producer + fan-in consumer). + dependentStageMap -= consumer + val events = deferral.delayedTaskCompletionEvents.toList + logInfo(log"Dropping ${MDC(NUM_EVENTS, events.size.toLong)} deferred completion(s) for " + + log"pipelined consumer ${MDC(STAGE, consumer)} because a producer failed; failing the " + + log"group") + // The buffered tasks genuinely succeeded, so still emit their TaskEnd events -- otherwise + // listeners that track active tasks (e.g. AppStatusListener) would believe these tasks are + // still running. We deliberately do NOT run the stage/job completion bookkeeping: the group + // is being failed as a unit, so the consumer's results must not be applied. + events.foreach(postTaskEnd) + } else if (deferral.parents.isEmpty) { + dependentStageMap -= consumer + val events = deferral.delayedTaskCompletionEvents.toList + logInfo(log"Replaying ${MDC(NUM_EVENTS, events.size.toLong)} deferred completion(s) " + + log"for pipelined consumer ${MDC(STAGE, consumer)} now that its producers finished") + events.foreach(eventProcessLoop.post) + } + } } /** @@ -3822,6 +4646,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 79f7af48f102..94078d698cf4 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,34 @@ 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. `createShuffleMapStage` still calls `MapOutputTracker.registerShuffle` + * for it (it is a `ShuffleDependency`, and an empty shuffle-status entry keeps the tracker's + * `containsShuffle` / `unregisterShuffle` bookkeeping uniform), but no map output is ever + * registered into that entry -- `registerMapOutput` runs only on the non-pipelined completion + * path -- so the entry stays empty. Its map-stage availability therefore cannot be read from the + * tracker; 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 +109,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 +123,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/TaskSchedulerImpl.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala index 618c8eb45902..196ae04d8acd 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala @@ -160,6 +160,37 @@ private[spark] class TaskSchedulerImpl( executorIdToRunningTaskIds.toMap.transform((_, v) => v.size) } + /** + * The number of outstanding tasks for the given resource profile that belong to work OTHER than + * the stages in `excludeStageIds`. "Outstanding" is the not-yet-completed demand of each task set + * -- running plus enqueued (`numTasks - tasksSuccessful`) -- not just the tasks actively running, + * so a neighbor's queued backlog is charged against capacity too. Used by the pipelined-group + * slot admission check (see DAGScheduler): it compares the group's demand against the slots left + * free after accounting for everything else in the SAME resource profile, so it excludes the + * group's own members (whose tasks would otherwise be charged against the group's own admission). + * + * Computed from the TaskSetManagers so it is resource-profile-scoped, skips zombie (superseded) + * attempts so a retried stage is not double-counted, and is taken under a single lock so the + * count is one consistent snapshot (both the per-profile total and the excluded members are read + * together). + */ + private[scheduler] def outstandingTasksForOtherWorkInProfile( + resourceProfileId: Int, excludeStageIds: Set[Int]): Int = synchronized { + taskSetsByStageIdAndAttempt.iterator.flatMap { case (stageId, attempts) => + if (excludeStageIds.contains(stageId)) { + Iterator.empty + } else { + attempts.valuesIterator + // Skip zombie attempts (superseded by a retry/kill): a stage can have both a zombie and a + // live attempt in this map at once, and the live attempt already re-runs the zombie's + // outstanding tasks -- counting both would double-count that stage's demand. Matches the + // !isZombie filtering used elsewhere on this map. + .filter(tsm => !tsm.isZombie && tsm.taskSet.resourceProfileId == resourceProfileId) + .map(tsm => math.max(0, tsm.numTasks - tsm.tasksSuccessful)) + } + }.sum + } + // The set of executors we have on each host; this is used to compute hostsAlive, which // in turn is used to decide when we can attain data locality on a given host protected val hostToExecutors = new HashMap[String, HashSet[String]] 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 3513cb1f9376..407ebc53fcff 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 78e4fc78b880..3237a7318564 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala @@ -69,6 +69,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 @@ -1105,16 +1114,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 } } @@ -1190,7 +1214,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/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index e12348e1be2d..3474478dba61 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 @@ -380,6 +380,20 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti shuffleMergeRegister: Boolean = true ) extends DAGScheduler( sc, taskScheduler, listenerBus, mapOutputTracker, blockManagerMaster, env, clock) { + + // Capacity reported to the pipelined-group slot check. Defaults high so scheduling-logic tests + // are not gated by the real (local[2]) backend capacity; a test can lower it to exercise the + // fail-fast path. + @volatile var maxConcurrentTasksForTest: Int = 1000 + override protected def maxConcurrentTasksForProfile(rpId: Int): Int = maxConcurrentTasksForTest + + // Seam for the free-slot admission check: outstanding (running + enqueued) task + // demand of OTHER work. Default 0 so existing tests see a fully-free cluster; a test can set it + // to model a busy neighbor. + @volatile var outstandingTasksForOtherWorkForTest: (Int, Set[Int]) => Int = (_, _) => 0 + override protected def outstandingTasksForOtherWork(rpId: Int, excludeStageIds: Set[Int]): Int = + outstandingTasksForOtherWorkForTest(rpId, excludeStageIds) + /** * Schedules shuffle merge finalize. */ @@ -6116,6 +6130,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assert(scheduler.runningStages.isEmpty) assert(scheduler.shuffleIdToMapStage.isEmpty) assert(scheduler.waitingStages.isEmpty) + assert(scheduler.dependentStageMap.isEmpty) assert(scheduler.outputCommitCoordinator.isEmpty) } @@ -6148,6 +6163,1569 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } CompletionEvent(task, reason, result, accumUpdates ++ extraAccumUpdates, metricPeaks, taskInfo) } + + // ========================================================================================== + // Pipelined shuffle dependency: group formation + concurrent submission + // ========================================================================================== + + test("pipelined shuffle: consumer stage is submitted concurrently with its producer") { + // producer (shuffle map) --[pipelined]--> consumer (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) + submit(consumerRdd, Array(0, 1)) + + // Both the producer's and the consumer's task sets must be submitted before the producer + // completes: the pipelined edge is non-sequencing, so the consumer does not wait. + assert(taskSets.size === 2, + s"expected producer and consumer submitted concurrently, got ${taskSets.size} task sets") + val producerStageId = taskSets.head.stageId + val consumerStageId = taskSets(1).stageId + assert(producerStageId != consumerStageId) + // The consumer is actually running (not parked in waitingStages), co-resident with the + // producer -- both are in runningStages. + assert(!scheduler.waitingStages.exists(_.rdd eq consumerRdd), + "consumer should not be waiting; it is co-scheduled with its pipelined producer") + assert(scheduler.runningStages.exists(_.rdd eq consumerRdd)) + assert(scheduler.runningStages.exists(_.rdd eq producerRdd)) + + // Now complete the producer, then the consumer, and the job finishes. + completeShuffleMapStageSuccessfully(producerStageId, 0, 2) + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("regular shuffle still waits (pipelined change is inert without a pipelined dependency)") { + // Identical shape but a REGULAR shuffle dependency: the consumer must NOT be submitted until + // the producer materializes -- verifying the concurrent-submission path is inert here. + 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)) + + // Only the producer is submitted initially; the consumer waits. + assert(taskSets.size === 1, s"expected only the producer submitted, got ${taskSets.size}") + val producerStageId = taskSets.head.stageId + completeShuffleMapStageSuccessfully(producerStageId, 0, 2) + // Now the consumer is submitted. + assert(taskSets.size === 2) + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a job mixing a pipelined and a regular shuffle is rejected up front") { + // A job must be either all-regular or all-pipelined, not a mix. A consumer + // depending on BOTH a pipelined producer AND a regular producer is a mixed job and must be + // rejected up front (before any stage is submitted), leaving no scheduler state behind. + val pipelinedProducerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(pipelinedProducerRdd, new HashPartitioner(2)) + val regularProducerRdd = new MyRDD(sc, 2, Nil) + val regularDep = new ShuffleDependency(regularProducerRdd, new HashPartitioner(2)) + val consumerRdd = + new MyRDD(sc, 2, List(pipelinedDep, regularDep), 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 mixed pipelined+regular job 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: 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) + val psdAB = new PipelinedShuffleDependency(rddA, new HashPartitioner(2)) + val rddB = new MyRDD(sc, 2, List(psdAB), tracker = mapOutputTracker) + val psdBC = new PipelinedShuffleDependency(rddB, new HashPartitioner(2)) + val rddC = new MyRDD(sc, 2, List(psdBC), tracker = mapOutputTracker) + submit(rddC, Array(0, 1)) + + // All three stages must be running concurrently; none parked. + assert(taskSets.size === 3, s"expected A, B, C co-scheduled, got ${taskSets.size}") + assert(scheduler.waitingStages.isEmpty, "no stage should be waiting in an all-pipelined chain") + Seq(rddA, rddB, rddC).foreach { rdd => + assert(scheduler.runningStages.exists(_.rdd eq rdd), s"stage for $rdd should be running") + } + + // Drain in dependency order: A, then B, then the result stage C. + val idA = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddA).get.stageId + completeShuffleMapStageSuccessfully(idA, 0, 2) + val idB = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddB).get.stageId + completeShuffleMapStageSuccessfully(idB, 0, 2) + val tsC = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddC).get + complete(tsC, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + 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. + sc.conf.set(config.SPECULATION_ENABLED, true) + try { + 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)) + // Not rejected: the producer is submitted and the job proceeds normally. + assert(taskSets.nonEmpty, "regular-shuffle job must not be rejected under speculation") + completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2) + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } finally { + sc.conf.set(config.SPECULATION_ENABLED, false) + } + } + + test("pipelined shuffle: speculation is rejected for a job with a pipelined dependency") { + // Speculation races a producer copy against a consumer reading its partial output; reject. + // The scheduler reads sc.conf live, so toggling it here takes effect for the submit below. + sc.conf.set(config.SPECULATION_ENABLED, true) + 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 = {} + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(consumerRdd, Array(0, 1), listener = failListener) + assert(failure.get() != null, + "job with pipelined dependency + speculation should be rejected") + assert(failure.get().getMessage.contains("Speculative execution is not supported")) + // No task sets were submitted for the rejected job. + assert(taskSets.isEmpty) + assertDataStructuresEmpty() + } finally { + sc.conf.set(config.SPECULATION_ENABLED, false) + } + } + + test("pipelined shuffle: dynamic allocation is rejected for a job with a pipelined dependency") { + // A pipelined group is gang-scheduled and its slot check measures currently-active executors; + // under dynamic allocation a job can start before any executor is up, so the check would fail + // the group against a transient 0-slot snapshot. Reject the combo (barrier does the same). + // DYN_ALLOCATION_TESTING lets isDynamicAllocationEnabled report true under a local master. + sc.conf.set(config.DYN_ALLOCATION_ENABLED, true) + sc.conf.set(config.DYN_ALLOCATION_TESTING, true) + 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 = {} + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(consumerRdd, Array(0, 1), listener = failListener) + assert(failure.get() != null, + "job with pipelined dependency + dynamic allocation should be rejected") + assert(failure.get().getMessage.contains("Dynamic allocation is not supported")) + assert(taskSets.isEmpty) + assertDataStructuresEmpty() + } finally { + sc.conf.set(config.DYN_ALLOCATION_ENABLED, false) + sc.conf.set(config.DYN_ALLOCATION_TESTING, false) + } + } + + 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. + sc.conf.set(config.DYN_ALLOCATION_ENABLED, true) + sc.conf.set(config.DYN_ALLOCATION_TESTING, true) + try { + val shuffleMapRdd = new MyRDD(sc, 2, Nil) + val shuffleDep = new ShuffleDependency(shuffleMapRdd, new HashPartitioner(2)) + val reduceRdd = new MyRDD(sc, 2, List(shuffleDep), tracker = mapOutputTracker) + submit(reduceRdd, Array(0, 1)) + assert(taskSets.nonEmpty, "regular-shuffle job must not be rejected under dynamic allocation") + } finally { + sc.conf.set(config.DYN_ALLOCATION_ENABLED, false) + sc.conf.set(config.DYN_ALLOCATION_TESTING, false) + } + } + + test("pipelined shuffle: submitting a pipelined dependency as a map-stage job is rejected") { + // A map-stage job (SparkContext.submitMapStage, e.g. AQE's ShuffleExchangeExec) materializes a + // shuffle to produce map-output statistics. A PipelinedShuffleDependency cannot serve that -- + // it + // is transient with no durable map output (getStatistics would return all-zero stats), and its + // producer/consumer co-scheduling makes speculation unsafe. It is rejected outright, regardless + // of speculation. No partial scheduler state is left behind. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + 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) + } + submitMapStage(pipelinedDep, listener = failListener) + assert(failure.get() != null, + "submitMapStage with a pipelined dependency must be rejected") + assert(failure.get().getMessage.contains("cannot be submitted as a map-stage job"), + s"expected a map-stage-rejection error, got: ${failure.get().getMessage}") + // No task sets were submitted for the rejected map-stage job, and no partial state remains. + assert(taskSets.isEmpty) + assertDataStructuresEmpty() + } + + test("regular shuffle map-stage job with speculation enabled is NOT rejected (inertness)") { + // The map-stage speculation fail-fast must be inert for a regular ShuffleDependency. + sc.conf.set(config.SPECULATION_ENABLED, true) + try { + val producerRdd = new MyRDD(sc, 2, Nil) + val regularDep = new ShuffleDependency(producerRdd, new HashPartitioner(2)) + submitMapStage(regularDep) + // Not rejected: the map stage is submitted and runs normally. + assert(taskSets.nonEmpty, + "regular-shuffle map-stage job must not be rejected under speculation") + completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2) + assertDataStructuresEmpty() + } finally { + sc.conf.set(config.SPECULATION_ENABLED, false) + } + } + + // ========================================================================================== + // Gang admission / slot check + // ========================================================================================== + + test("pipelined shuffle: an all-pipelined group that fits is admitted up front and runs") { + // Whole-group demand producer(2) + consumer(2) = 4 <= capacity 4, other-work occupancy 0, so + // the up-front gang admission check (handleJobSubmitted) admits the job before any stage runs, + // and both stages are then co-scheduled and complete normally. + val producerRdd = new MyRDD(sc, 2, Nil) + val myScheduler = scheduler.asInstanceOf[MyDAGScheduler] + myScheduler.maxConcurrentTasksForTest = 4 + myScheduler.outstandingTasksForOtherWorkForTest = (_, _) => 0 + try { + 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, "a group that fits must be admitted and co-scheduled") + completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2) + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } finally { + myScheduler.maxConcurrentTasksForTest = 1000 + myScheduler.outstandingTasksForOtherWorkForTest = (_, _) => 0 + } + } + + test("pipelined shuffle: a group too large to co-fit fails fast with INSUFFICIENT_SLOT") { + // Constrain reported capacity to 3 slots. A pipelined group of producer(2) + consumer(2) = 4 + // tasks exceeds it and can never be co-resident, so it must fail fast rather than deadlock. + val producerRdd = new MyRDD(sc, 2, Nil) // touch sc first so `scheduler` is initialized + scheduler.asInstanceOf[MyDAGScheduler].maxConcurrentTasksForTest = 3 + try { + 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, "an over-large pipelined group must fail the job") + assert(failure.get().getMessage.contains("CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT") || + failure.get().getMessage.contains("concurrent task slots"), + s"expected an insufficient-slot error, got: ${failure.get().getMessage}") + assertDataStructuresEmpty() + } finally { + scheduler.asInstanceOf[MyDAGScheduler].maxConcurrentTasksForTest = 1000 + } + } + + test("pipelined shuffle: an over-capacity 3-stage all-pipelined chain is rejected up front") { + // A 3-stage all-pipelined chain A->B->C has whole-group demand 2+2+2 = 6. With capacity 3 it + // cannot co-fit, so the up-front gang admission check (handleJobSubmitted) fails the job before + // any stage is created -- leaving no scheduler state behind (assertDataStructuresEmpty). This + // covers a multi-stage group in addition to the 2-stage "group too large" case. + val rddA = new MyRDD(sc, 2, Nil) // touch sc first so `scheduler` is initialized + scheduler.asInstanceOf[MyDAGScheduler].maxConcurrentTasksForTest = 3 + try { + val depAB = new PipelinedShuffleDependency(rddA, new HashPartitioner(2)) + val rddB = new MyRDD(sc, 2, List(depAB), tracker = mapOutputTracker) + val depBC = new PipelinedShuffleDependency(rddB, new HashPartitioner(2)) + val rddC = new MyRDD(sc, 2, List(depBC), 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(rddC, Array(0, 1), listener = failListener) + + assert(failure.get() != null, "an over-large 3-stage pipelined chain must fail the job") + assert(failure.get().getMessage.contains("CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT") || + failure.get().getMessage.contains("concurrent task slots"), + s"expected an insufficient-slot error, got: ${failure.get().getMessage}") + // No stage is created for a job rejected up front. + assert(taskSets.isEmpty, "no stage should be submitted for an up-front-rejected group") + assertDataStructuresEmpty() + } finally { + scheduler.asInstanceOf[MyDAGScheduler].maxConcurrentTasksForTest = 1000 + } + } + + test("pipelined shuffle: a group that fits is co-scheduled (slot check passes)") { + // Producer (2) + consumer (2) = 4 tasks <= 16 slots: co-scheduled normally, no abort. + 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, "group fits, so producer and consumer are co-scheduled") + completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2) + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a group that fits total capacity but not FREE slots fails fast") { + // Admission is decided against currently-FREE slots, not total capacity. A group of + // producer(2) + consumer(2) = 4 tasks fits a 10-slot cluster in principle, but if 8 slots are + // already occupied by OTHER work only 2 are free -- the group cannot co-fit right now and must + // fail fast rather than queue forever. Under the old total-capacity check this would wrongly be + // admitted (4 <= 10) and then hang. outstandingTasksForOtherWork already excludes the group's + // own members, so we report 8 directly. + val producerRdd = new MyRDD(sc, 2, Nil) // touch sc first so `scheduler` is initialized + val myScheduler = scheduler.asInstanceOf[MyDAGScheduler] + myScheduler.maxConcurrentTasksForTest = 10 + // 8 of 10 slots busy elsewhere -> 2 free + myScheduler.outstandingTasksForOtherWorkForTest = (_, _) => 8 + try { + 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 group that fits total but not free slots must fail the job (free-slot check)") + assert(failure.get().getMessage.contains("CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT") || + failure.get().getMessage.contains("currently free"), + s"expected a free-slot insufficient-slot error, got: ${failure.get().getMessage}") + assertDataStructuresEmpty() + } finally { + myScheduler.maxConcurrentTasksForTest = 1000 + myScheduler.outstandingTasksForOtherWorkForTest = (_, _) => 0 + } + } + + test("pipelined shuffle: slot check can be disabled, admitting a group that exceeds capacity") { + // With spark.scheduler.pipelinedGroup.slotCheck.enabled=false the admission check is skipped + // entirely, so a group whose demand exceeds free slots is co-scheduled instead of failing fast. + // Deployments that admit capacity out-of-band (e.g. a slot reservation) rely on this. + val producerRdd = new MyRDD(sc, 2, Nil) + val myScheduler = scheduler.asInstanceOf[MyDAGScheduler] + myScheduler.maxConcurrentTasksForTest = 2 // demand 4 > 2 slots: would fail if the check ran + myScheduler.outstandingTasksForOtherWorkForTest = (_, _) => 0 + val prev = sc.conf.get(config.PIPELINED_GROUP_SLOT_CHECK_ENABLED) + sc.conf.set(config.PIPELINED_GROUP_SLOT_CHECK_ENABLED, false) + try { + 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, + "with the slot check disabled, the over-capacity group is co-scheduled, not rejected") + completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2) + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } finally { + myScheduler.maxConcurrentTasksForTest = 1000 + sc.conf.set(config.PIPELINED_GROUP_SLOT_CHECK_ENABLED, prev) + } + } + + test("pipelined shuffle: a group whose demand exactly equals free capacity is admitted") { + // Boundary: demand == free capacity must be admitted, not rejected by an off-by-one (the check + // is `demand > freeSlots`, not `>=`). 4 total slots, other-work occupancy 0 -> free 4 >= demand + // 4 -> admitted. NOTE: the DAGScheduler up-front admission (rejectUnadmittablePipelinedGroup) + // passes excludeStageIds = Set.empty (no job stage exists yet at admission time), so the + // group's-own-members exclusion is NOT exercised here; that exclusion is genuinely covered by + // TaskSchedulerImplSuite ("outstandingTasksForOtherWorkInProfile excludes the given stages"). + val producerRdd = new MyRDD(sc, 2, Nil) + val myScheduler = scheduler.asInstanceOf[MyDAGScheduler] + myScheduler.maxConcurrentTasksForTest = 4 + myScheduler.outstandingTasksForOtherWorkForTest = (_, _) => 0 + try { + 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, + "a group whose demand equals free capacity must be co-scheduled, not rejected") + completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2) + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } finally { + myScheduler.maxConcurrentTasksForTest = 1000 + myScheduler.outstandingTasksForOtherWorkForTest = (_, _) => 0 + } + } + + // ========================================================================================== + // Group-observable completion + // ========================================================================================== + + test("pipelined shuffle: a consumer that finishes early does not end the job or cancel its " + + "still-running producer") { + // producer (shuffle map) --[pipelined]--> consumer (result). Consumer is co-scheduled; if its + // result tasks finish before the producer, its completion MUST be deferred -- otherwise the job + // would end and cancelRunningIndependentStages would cancel the still-running 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)) + assert(taskSets.size === 2) + val producerStageId = taskSets.head.stageId + val consumerTaskSet = taskSets(1) + + // Consumer finishes FIRST (both result tasks succeed) while the producer is still running. + complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) + // The job must NOT be finished yet, and the producer must still be running (not cancelled). + assert(results.isEmpty, "consumer completion must be deferred until the producer finishes") + assert(scheduler.runningStages.exists(_.rdd eq producerRdd), + "the still-running producer must not be cancelled by the early-finishing consumer") + assert(cancelledStages.isEmpty) + + // Now the producer finishes -> the deferred consumer completions replay -> the job completes. + completeShuffleMapStageSuccessfully(producerStageId, 0, 2) + assert(results === Map(0 -> 42, 1 -> 43), + "deferred consumer completions should replay once the producer finishes") + assertDataStructuresEmpty() + } + + test("pipelined shuffle: producer-then-consumer ordering completes normally (no deferral " + + "needed)") { + // Sanity: when the producer finishes before the consumer's tasks complete, there is nothing to + // defer and the job completes as usual. + 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 consumerTaskSet = taskSets(1) + + completeShuffleMapStageSuccessfully(producerStageId, 0, 2) + complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: deferred consumer completions are dropped when the producer fails") { + // If a producer fails while a co-scheduled consumer's completions are buffered, those buffered + // successes must be DROPPED (not applied): the group is torn down / the job fails, and applying + // a partial consumer success would be incorrect (S6). The buffered tasks genuinely succeeded, + // so their TaskEnd events are still emitted on the drop path (otherwise listeners that track + // active tasks would believe these tasks are still running); only the stage/job-completion + // bookkeeping is withheld. + val taskEndCount = new java.util.concurrent.atomic.AtomicInteger(0) + val countingListener = new SparkListener { + override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = taskEndCount.incrementAndGet() + } + sc.addSparkListener(countingListener) + 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) + val producerTaskSet = taskSets.head + val consumerTaskSet = taskSets(1) + + // Consumer finishes early -> its whole completion events are buffered (deferred): no result + // and no TaskEnd yet, since the entire event is withheld until the producer's fate is known. + complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) + assert(results.isEmpty, "consumer completions should be buffered while the producer runs") + sc.listenerBus.waitUntilEmpty(10000) + assert(taskEndCount.get() === 0, + "buffered consumer completions must not fire TaskEnd until released; got " + + taskEndCount.get()) + + // The producer now FAILS its whole task set. The job fails; the buffered consumer successes + // must NOT be applied as results, but their TaskEnd events ARE emitted on the drop path so + // active-task-tracking listeners see the tasks finish. + failed(producerTaskSet, "producer blew up") + assert(failure.get() != null, "job should fail when the producer fails") + assert(results.isEmpty, + "buffered consumer successes must be dropped when the producer fails, not applied") + sc.listenerBus.waitUntilEmpty(10000) + assert(taskEndCount.get() === 2, + "drop path must emit the buffered consumer TaskEnd events; got " + taskEndCount.get()) + assertDataStructuresEmpty() + } finally { + sc.removeSparkListener(countingListener) + } + } + + test("pipelined shuffle: a fan-in consumer's buffered successes are dropped as soon as one " + + "producer fails, even if another later succeeds") { + // Fan-in: a consumer co-scheduled with TWO pipelined producers (A and B). If A fails, the + // consumer's buffered successes must be DROPPED -- they were computed against A's now-invalid + // output, and the group is failed as a unit (S6). The drop happens immediately when A fails, + // not deferred until the last producer finishes: so even if B later succeeds, the consumer has + // already been dropped and B's completion is a clean no-op (it must NOT replay the successes). + val producerA = new MyRDD(sc, 2, Nil) + val psdA = new PipelinedShuffleDependency(producerA, new HashPartitioner(2)) + val producerB = new MyRDD(sc, 2, Nil) + val psdB = new PipelinedShuffleDependency(producerB, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(psdA, psdB), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + assert(taskSets.size === 3, s"A, B and the consumer must be co-scheduled; got ${taskSets.size}") + val consumerTaskSet = + taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd).get + val stageA = scheduler.stageIdToStage( + taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq producerA).get.stageId) + val stageB = scheduler.stageIdToStage( + taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq producerB).get.stageId) + val consumerStage = scheduler.stageIdToStage(consumerTaskSet.stageId) + + // Consumer finishes early -> its successes are buffered against BOTH A and B. + complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) + assert(results.isEmpty, "consumer completions should be buffered while its producers run") + assert(scheduler.dependentStageMap.get(consumerStage).exists(_.parents.size == 2), + "the consumer must be deferred against both producers") + + // A fails: the consumer is dropped immediately (not retained pending B), and no result applied. + // (Driven via markStageAsFinished directly; the job-failure teardown that would normally follow + // is exercised by the group-atomic-failure tests in a later PR -- here we assert only the + // drop-vs-replay outcome.) + scheduler.markStageAsFinished(stageA, errorMessage = Some("producer A blew up")) + assert(!scheduler.dependentStageMap.contains(consumerStage), + "the consumer's deferral must be dropped as soon as a producer fails") + assert(results.isEmpty, "a fan-in consumer's buffered successes must not be applied on failure") + + // B now succeeds: the consumer is already gone, so this is a clean no-op -- it must NOT replay. + scheduler.runningStages += stageB + completeShuffleMapStageSuccessfully(stageB.id, 0, 2) + assert(results.isEmpty, + 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 + // update) applied exactly once -- at replay -- not once when buffered and again when replayed. + // The producer (2 tasks) and consumer (2 tasks) yield exactly 4 TaskEnd events total; a + // buffer+replay double-post would inflate that count. + val taskEndCount = new java.util.concurrent.atomic.AtomicInteger(0) + val countingListener = new SparkListener { + override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = taskEndCount.incrementAndGet() + } + sc.addSparkListener(countingListener) + 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) + submit(consumerRdd, Array(0, 1)) + val producerStageId = taskSets.head.stageId + val consumerTaskSet = taskSets(1) + + // Consumer finishes first -> its two completions are buffered (deferred, no TaskEnd yet). + complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) + assert(results.isEmpty, "consumer job result must be deferred until the producer finishes") + sc.listenerBus.waitUntilEmpty(10000) + assert(taskEndCount.get() === 0, + "deferred consumer completions must not fire TaskEnd until replay; got " + + taskEndCount.get()) + + // Producer finishes -> the two deferred consumer completions replay. + completeShuffleMapStageSuccessfully(producerStageId, 0, 2) + assert(results === Map(0 -> 42, 1 -> 43)) + sc.listenerBus.waitUntilEmpty(10000) + + // Exactly 4 TaskEnd events (2 producer + 2 consumer). A double-post from buffer+replay would + // make it 6 (the 2 consumer tasks counted twice). + assert(taskEndCount.get() === 4, + s"expected 4 TaskEnd events (no buffer+replay duplication), got ${taskEndCount.get()}") + assertDataStructuresEmpty() + } finally { + sc.removeSparkListener(countingListener) + } + } + + test("pipelined shuffle: releasing a deferral is gated on the producer being genuinely done") { + // Regression for the "producer finished but not available -> resubmitted" hazard: a + // ShuffleMapStage producer that reaches markStageAsFinished with no error but is NOT available + // (some output missing, so it will be resubmitted) must NOT release/replay its deferred + // consumers -- doing so would apply the consumer's results against soon-to-be-recomputed + // output. + // This drives BOTH branches of the gate: (1) the RETAINED branch -- markStageAsFinished on a + // not-yet-available pipelined producer must leave the deferral in place; and (2) the RELEASED + // branch -- once the producer is genuinely available, the deferral is released and results + // applied. + 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 consumerTaskSet = taskSets(1) + + complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) + // Deferred while the producer runs. + assert(scheduler.dependentStageMap.keys.exists(_.rdd eq consumerRdd)) + assert(results.isEmpty) + + val producerStage = + scheduler.stageIdToStage(producerStageId).asInstanceOf[ShuffleMapStage] + + // (1) RETAINED branch: the producer has completed NO partitions yet, so it is not available. + // markStageAsFinished with no error and willRetry=false hits the producerAboutToResubmit guard + // (ShuffleMapStage && !producerFailed && !isAvailable) and must NOT release the consumer. + assert(!producerStage.isAvailable, "precondition: producer not yet available") + scheduler.markStageAsFinished(producerStage, errorMessage = None, willRetry = false) + assert(scheduler.dependentStageMap.keys.exists(_.rdd eq consumerRdd), + "a not-yet-available pipelined producer must NOT release its deferred consumers (it is " + + "about to resubmit; releasing would apply results against soon-to-be-recomputed output)") + assert(results.isEmpty, + "consumer results must not be applied while the producer is not available") + + // (2) RELEASED branch: producer completes for real and IS available -> deferral released, + // consumer results applied. (Re-add it to runningStages since (1) removed it, so the normal + // completion path proceeds as in production.) + scheduler.runningStages += producerStage + completeShuffleMapStageSuccessfully(producerStageId, 0, 2) + assert(producerStage.isAvailable) + assert(!scheduler.dependentStageMap.keys.exists(_.rdd eq consumerRdd), + "deferral released once the producer is genuinely done (available)") + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: job teardown flushes a buffered consumer's TaskEnds even when the " + + "release path never drained it") { + // A consumer's whole completion event is buffered while its producer runs (coarse model), so + // its TaskEnds are held, not yet emitted. Normally releaseDeferredPipelinedConsumers drains + // them (cancelRunningIndependentStages finishes each running producer before cleanup). But a + // producer left in resubmit limbo -- finished with no error yet NOT available, so + // producerAboutToResubmit held the deferral AND markStageAsFinished removed it from + // runningStages -- is skipped by cancelRunningIndependentStages (neither running nor failed). + // If the job is then cancelled, cleanup must still flush the buffered TaskEnds (results stay + // dropped) so a listener tracking active tasks does not leak them as perpetually running. + val taskEndCount = new java.util.concurrent.atomic.AtomicInteger(0) + val countingListener = new SparkListener { + override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = taskEndCount.incrementAndGet() + } + sc.addSparkListener(countingListener) + 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 jobId = submit(consumerRdd, Array(0, 1)) + val producerStageId = taskSets.head.stageId + val consumerTaskSet = taskSets(1) + + // Consumer finishes early -> whole completion buffered; no TaskEnd yet. + complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) + sc.listenerBus.waitUntilEmpty() + assert(taskEndCount.get() === 0, "buffered consumer TaskEnds are held while producer runs") + + // Drive the producer into resubmit limbo: finished, no error, but not available -> the + // deferral is retained and the producer leaves runningStages. No task failed, so it also has + // no failedAttemptIds -- cancelRunningIndependentStages will skip it. + val producerStage = scheduler.stageIdToStage(producerStageId).asInstanceOf[ShuffleMapStage] + assert(!producerStage.isAvailable, "precondition: producer not available") + scheduler.markStageAsFinished(producerStage, errorMessage = None, willRetry = false) + assert(scheduler.dependentStageMap.keys.exists(_.rdd eq consumerRdd), + "precondition: the deferral is retained (producer about to resubmit)") + assert(!scheduler.runningStages.contains(producerStage) && + producerStage.failedAttemptIds.isEmpty, + "precondition: producer is neither running nor failed, so cancel will skip it") + + // Cancel the job. The release path never drains this consumer, so cleanup must flush it. + cancel(jobId) + sc.listenerBus.waitUntilEmpty() + assert(scheduler.dependentStageMap.isEmpty, "the deferral must not outlive the job") + assert(results.isEmpty, "a cancelled job's buffered consumer success must not be applied") + assert(taskEndCount.get() === 2, + s"teardown must flush the 2 buffered consumer TaskEnds (else active tasks leak); got " + + taskEndCount.get()) + assertDataStructuresEmpty() + } finally { + sc.removeSparkListener(countingListener) + } + } + + 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 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/TaskSchedulerImplSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala index c9c9e529405e..edf2dcdc0e75 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala @@ -2752,4 +2752,62 @@ class TaskSchedulerImplSuite extends SparkFunSuite with LocalSparkContext assert(!tsm.isInstanceOf[StructuredStreamingIdAwareSchedulerLogging]) } + test("outstandingTasksForOtherWorkInProfile counts running + enqueued, excludes given stages, " + + "and is resource-profile-scoped") { + // Backs the pipelined-group slot admission check (DAGScheduler): it must report the OUTSTANDING + // (running + enqueued, i.e. numTasks - tasksSuccessful) task demand of OTHER work in a given + // resource profile, excluding the group's own member stages. + val taskScheduler = setupScheduler() + // Only 3 total cores, but each stage has 4 tasks -> some tasks launch, the rest stay enqueued. + val workerOffers = IndexedSeq( + new WorkerOffer("executor0", "host0", 2), + new WorkerOffer("executor1", "host1", 1)) + + // Stage 0 and stage 1 both in the DEFAULT profile, 4 tasks each (8 total) but only 3 cores, so + // 3 launch and 5 remain enqueued. Distinct stageIds let us exclude one and count the other. + taskScheduler.submitTasks(FakeTask.createTaskSet(4, stageId = 0, stageAttemptId = 0)) + taskScheduler.submitTasks(FakeTask.createTaskSet(4, stageId = 1, stageAttemptId = 0)) + val launched = taskScheduler.resourceOffers(workerOffers).flatten + assert(launched.length === 3, "only three tasks can run (3 cores); the other five are enqueued") + + val defaultRp = ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID + // Nothing excluded: all 8 tasks are outstanding (3 running + 5 enqueued) though only 3 run. + // The old running-only count would have reported 3 here; outstanding demand is what can hang a + // co-scheduled group, so it must be 8. + assert(taskScheduler.outstandingTasksForOtherWorkInProfile(defaultRp, Set.empty) === 8) + // Exclude stage 0 (as if it were the group's own member): only stage 1's 4 tasks remain. + assert(taskScheduler.outstandingTasksForOtherWorkInProfile(defaultRp, Set(0)) === 4) + // Exclude both: zero "other" tasks. + assert(taskScheduler.outstandingTasksForOtherWorkInProfile(defaultRp, Set(0, 1)) === 0) + // A DIFFERENT (non-existent here) resource profile has no outstanding tasks -- other-profile + // load must not be charged against this profile (RP-scoping). + val otherRpId = defaultRp + 12345 + assert(taskScheduler.outstandingTasksForOtherWorkInProfile(otherRpId, Set.empty) === 0, + "tasks in the default profile must not be counted against a different profile") + } + + test("outstandingTasksForOtherWorkInProfile does not double-count a zombie + live attempt") { + // A retried/killed stage can have both a zombie (superseded) attempt and a live attempt in the + // map at once; the live attempt re-runs the zombie's outstanding tasks, so only the live + // attempt's demand should count. Counting both would inflate occupancy and could spuriously + // fail a pipelined group with CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT. + val taskScheduler = setupScheduler() + val defaultRp = ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID + + // Attempt 0 of stage 0: 4 tasks, none started. Mark it zombie (as if superseded). + val attempt0 = FakeTask.createTaskSet(4, stageId = 0, stageAttemptId = 0) + taskScheduler.submitTasks(attempt0) + taskScheduler.taskSetManagerForAttempt(0, 0).get.isZombie = true + + // Attempt 1 of the SAME stage: the live retry, also 4 tasks. + val attempt1 = FakeTask.createTaskSet(4, stageId = 0, stageAttemptId = 1) + taskScheduler.submitTasks(attempt1) + + // Both attempts are in the map, but only the live attempt's 4 tasks are outstanding -- not 8. + assert(taskScheduler.outstandingTasksForOtherWorkInProfile(defaultRp, Set.empty) === 4, + "a zombie attempt must not be counted alongside its live retry") + // Excluding the stage (as the group's own member) drops both attempts. + assert(taskScheduler.outstandingTasksForOtherWorkInProfile(defaultRp, Set(0)) === 0) + } + } 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 28d774eb27f7..4c7e8efc4eab 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) {