You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by we...@apache.org on 2026/07/30 12:16:53 UTC
(spark) branch master updated: [SPARK-58398][CORE] Group-atomic failure and fail-fast rejection for pipelined-shuffle stage groups
This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 188544238d61 [SPARK-58398][CORE] Group-atomic failure and fail-fast rejection for pipelined-shuffle stage groups
188544238d61 is described below
commit 188544238d613c2bb5ddfe982efcfc93bf94e36f
Author: Boyang Jerry Peng <je...@databricks.com>
AuthorDate: Thu Jul 30 20:16:33 2026 +0800
[SPARK-58398][CORE] Group-atomic failure and fail-fast rejection for pipelined-shuffle stage groups
### What changes were proposed in this pull request?
This PR combines the two failure-handling layers of the pipelined-shuffle work into a single
reviewable change, on top of the now-merged concurrent scheduling of a `PipelinedShuffleDependency`
group (SPARK-58263 / #57341): the **group-atomic failure model** and the **fail-fast rejection of
unsupported idioms** that together make co-scheduling safe under failure. (It builds on two merged
predecessors -- the `PipelinedShuffleDependency` definition and type-based routing from
SPARK-58185 / #57286, and the concurrent scheduling from SPARK-58263 / #57341 -- and targets
`master` directly; its net-new content is the failure and fail-fast layers described below.)
A pipelined shuffle is transient and incrementally read: its output is streamed to a co-scheduled
consumer and is never materialized as durable, re-readable map output. That has two consequences the
scheduler must enforce.
**Group-atomic failure (a pipelined group succeeds or fails as a whole).** A transient shuffle
cannot be recomputed in isolation and its members run concurrently, so the stock "resubmit one
stage" recovery does not apply -- any member failure must fail the whole group, which the caller
then reruns (a streaming query restarts the batch).
- **Member task failure.** A pipelined group member's `TaskSet` is tagged `isPipelined` and pinned
to `maxTaskFailures = 1`, so the first task failure aborts the set instead of retrying in place,
which routes to a whole-group abort.
- **Executor loss.** Losing an executor running a member task is force-counted as a member failure
(except benign `TaskKilled` / `TaskCommitDenied`), aborting the group.
- **FetchFailed on a member.** Handled by a dedicated branch that aborts the whole group rather than
resubmitting the map stage in isolation (a lone-stage resubmit of a transient shuffle would
deadlock the group). The failed executor's *regular* outputs are still unregistered from the
`MapOutputTracker`, exactly as the base handler does, for the benefit of other jobs.
- **No transient-producer resubmit.** A pipelined `ShuffleMapStage` records its completed partitions
locally and monotonically (never in the `MapOutputTracker`) and never flips back to unavailable, so
losing an executor that held an already-consumed pipelined output cannot make the scheduler
resubmit the producer -- which would otherwise hang the producer's streaming writer waiting on
termination acks from reducers that already finished. The `TaskSetManager` "Resubmitted"
re-enqueue loop likewise excludes pipelined sets.
- **Cross-job / cross-time reuse rejected.** A transient shuffle has no retained output for a second
job to read, so binding a pipelined producer stage to a second job fails fast.
**Fail-fast on unsupported idioms (spec S9).** A pipelined group is rejected up front, before any
stage is created (so a rejection leaves no partial scheduler state), when it uses an idiom v1 cannot
support: a producer feeding **more than one consumer** (1:N fan-out needs multicast, deferred); a
**barrier**, **statically-indeterminate**, **checksum-mismatch-retry**, or **push-merge** producer;
a **reliable RDD checkpoint** anywhere in a member's within-stage chain (it reintroduces cross-time
reuse of a transient edge); or **members with differing resource profiles** (the gang slot check
compares one demand against one profile's capacity, so v1 requires a single-profile group). These
throw a typed `PipelinedShuffleUnsupportedException` (carrying the `PIPELINED_SHUFFLE_UNSUPPORTED`
error class), which `handleJobSubmitted` matches by type.
Main changes:
- `DAGScheduler.scala` -- the FetchFailed group-abort branch, the no-resubmit handling of a
pipelined `ShuffleMapStage`, cross-job reuse rejection, and `checkPipelinedGroupsSupportedInRDDGraph`
/ `checkPipelinedProducerSupported` fail-fast (typed exception).
- `TaskSetManager.scala` -- `maxTaskFailures = 1` and force-counted executor loss for a pipelined
set; exclusion from the "Resubmitted" re-enqueue loop.
- `ShuffleMapStage.scala` -- monotonic local availability for a pipelined shuffle.
- `PIPELINED_SHUFFLE_UNSUPPORTED` and `PIPELINED_SHUFFLE_CROSS_JOB_REUSE` error conditions.
### Why are the changes needed?
Co-scheduling a pipelined group (#57341) is only safe if failure is handled at the granularity of
the whole group: because the shuffle is transient and once-through, the stock per-stage resubmit
recovery would either deadlock the group or hang a producer's streaming writer. This PR makes any
member failure fail the group atomically (so the caller reruns it) and rejects up front the idioms
whose recovery/semantics are incompatible with a transient, concurrently-read shuffle -- turning
what would be a hang or a silently-wrong schedule into a clear, immediate failure. The two layers
are combined into one PR because they are inseparable in review: the fail-fast rejections define
exactly which group shapes the failure model must then handle, and both are gated on the same
`PipelinedShuffleDependency` type.
### Does this PR introduce _any_ user-facing change?
No. All new behavior is gated on a job using a `PipelinedShuffleDependency`, which nothing constructs
yet, so every existing job is scheduled and recovers exactly as before. The new error conditions
(`PIPELINED_SHUFFLE_UNSUPPORTED`, `PIPELINED_SHUFFLE_CROSS_JOB_REUSE`) can only surface for a job
that uses a pipelined dependency.
### How was this patch tested?
New unit tests in `DAGSchedulerSuite` and `TaskSetManagerSuite` cover:
- group-atomic failure: `maxTaskFailures = 1` aborting a member set on the first failure; executor
loss force-counted (and benign `TaskKilled` / `TaskCommitDenied` not force-counted); a member
FetchFailed aborting the whole group rather than resubmitting a single stage;
- no transient-producer resubmit: a post-executor-loss straggler success not resubmitting the
producer; a completed pipelined producer's availability surviving executor loss; the "Resubmitted"
loop excluding a pipelined set (including the partial-producer-on-decommission case);
- cross-job reuse rejected; a group-atomic rerun resetting per-partition commit authorization;
- fail-fast idioms: fan-out, barrier / indeterminate / checksum / push-merge producer, a reliable
checkpoint in a producer's or a consumer's chain (including downstream in the consumer stage), and
a mixed-resource-profile group -- each rejected up front; and that regular-shuffle idioms are NOT
rejected (inertness of the fail-fast for a job with no pipelined dependency).
The full `DAGSchedulerSuite` and `TaskSetManagerSuite` pass.
### Was this patch authored or co-authored using generative AI tooling?
Co-authored with Claude Code (Opus 4.8)
Closes #57361 from jerrypeng/stack/pipelined-shuffle-pr6-failfast.
Authored-by: Boyang Jerry Peng <je...@databricks.com>
Signed-off-by: Wenchen Fan <we...@databricks.com>
---
.../src/main/resources/error/error-conditions.json | 12 +
.../org/apache/spark/scheduler/ActiveJob.scala | 9 +-
.../org/apache/spark/scheduler/DAGScheduler.scala | 456 +++++++---
.../apache/spark/scheduler/ShuffleMapStage.scala | 47 +-
.../scala/org/apache/spark/scheduler/TaskSet.scala | 8 +-
.../apache/spark/scheduler/TaskSetManager.scala | 46 +-
.../apache/spark/scheduler/DAGSchedulerSuite.scala | 972 +++++++++++++++++++--
.../spark/scheduler/TaskSetManagerSuite.scala | 198 +++++
8 files changed, 1555 insertions(+), 193 deletions(-)
diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json
index c7458b4cd3d6..d27b0c02b1d5 100644
--- a/common/utils/src/main/resources/error/error-conditions.json
+++ b/common/utils/src/main/resources/error/error-conditions.json
@@ -6438,6 +6438,18 @@
],
"sqlState" : "42K03"
},
+ "PIPELINED_SHUFFLE_CROSS_JOB_REUSE" : {
+ "message" : [
+ "The pipelined shuffle <shuffleId> is being reused across jobs. A pipelined (incrementally-readable) shuffle is transient and has no retained output for another job to read, so its producer stage cannot be shared across jobs."
+ ],
+ "sqlState" : "0A000"
+ },
+ "PIPELINED_SHUFFLE_UNSUPPORTED" : {
+ "message" : [
+ "A pipelined shuffle stage group cannot be scheduled because it uses an unsupported feature: <reason>. Pipelined (incrementally-readable) shuffles run their producer and consumer stages concurrently over a transient, once-through stream, which is incompatible with this feature in this version."
+ ],
+ "sqlState" : "0A000"
+ },
"PIPELINE_DATASET_WITHOUT_FLOW" : {
"message" : [
"Pipeline dataset <identifier> does not have any defined flows. Please attach a query with the dataset's definition, or explicitly define at least one flow that writes to the dataset."
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..0ea568693aeb 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/ActiveJob.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/ActiveJob.scala
@@ -48,7 +48,14 @@ private[spark] class ActiveJob(
val callSite: CallSite,
val listener: JobListener,
val artifacts: JobArtifactSet,
- val properties: Properties) {
+ val properties: Properties,
+ // Whether this job's RDD graph uses a `PipelinedShuffleDependency`. Every pipelined-group
+ // scheduling path -- co-scheduling, deferral, and the per-submit `TaskSet.isPipelined` tag --
+ // 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. Defaults to false; the
+ // result-job path (handleJobSubmitted) passes the value it computed up front, and the
+ // map-stage-job path (which rejects pipelined dependencies) leaves it false.
+ val hasPipelinedDependency: Boolean = false) {
/**
* Number of partitions we need to compute for this job. Note that result stages may not need
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 bfd24ceb2f89..a72af257ad73 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
@@ -45,7 +45,7 @@ import org.apache.spark.network.shuffle.{BlockStoreClient, MergeFinalizerListene
import org.apache.spark.network.shuffle.protocol.MergeStatuses
import org.apache.spark.network.util.JavaUtils
import org.apache.spark.partial.{ApproximateActionListener, ApproximateEvaluator, PartialResult}
-import org.apache.spark.rdd.{RDD, RDDCheckpointData}
+import org.apache.spark.rdd.{DeterministicLevel, RDD, RDDCheckpointData, ReliableRDDCheckpointData}
import org.apache.spark.resource.{CpuAmount, ResourceProfile, TaskResourceProfile}
import org.apache.spark.resource.ResourceProfile.{CPUS, DEFAULT_RESOURCE_PROFILE_ID, EXECUTOR_CORES_LOCAL_PROPERTY, MAX_TASKS_PER_EXECUTOR_LOCAL_PROPERTY, PYSPARK_MEMORY_LOCAL_PROPERTY}
import org.apache.spark.rpc.RpcTimeout
@@ -187,6 +187,11 @@ private[spark] class DAGScheduler(
delayedTaskCompletionEvents: ListBuffer[CompletionEvent] = new ListBuffer[CompletionEvent])
private[scheduler] val dependentStageMap = new HashMap[Stage, DependentStageInfo]
+ // Whether we have already logged that the pipelined-group slot check is disabled. Only the
+ // (single-threaded) event loop touches this, so a plain var is safe. Keeps the warning to once
+ // per scheduler rather than once per submitted batch job.
+ private var warnedPipelinedSlotCheckDisabled = false
+
private[scheduler] val activeJobs = new HashSet[ActiveJob]
// Track all the jobs submitted by the same query execution, will clean up after
@@ -622,6 +627,21 @@ private[spark] class DAGScheduler(
firstJobId: Int): ShuffleMapStage = {
shuffleIdToMapStage.get(shuffleDep.shuffleId) match {
case Some(stage) =>
+ // A pipelined shuffle is transient: it is a once-through live stream with no retained,
+ // addressable output, so reusing its producer stage across jobs is unsound (there is no
+ // durable output for a second job to read). Reuse must be prevented explicitly -- from the
+ // scheduler's view a shuffle-map stage can be reused unless something forbids it. If a
+ // pipelined dependency's shuffleId is already bound to a stage from a different job, that
+ // is the forbidden cross-job reuse; fail fast. (Within the same job the cached stage is the
+ // one we just created, so returning it is correct and not reuse.)
+ if (shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]] &&
+ !stage.jobIds.contains(firstJobId)) {
+ throw new SparkException(
+ errorClass = "PIPELINED_SHUFFLE_CROSS_JOB_REUSE",
+ messageParameters = scala.collection.immutable.Map(
+ "shuffleId" -> shuffleDep.shuffleId.toString),
+ cause = null)
+ }
stage
case None =>
@@ -671,6 +691,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()
@@ -694,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
@@ -1060,6 +1138,126 @@ private[spark] class DAGScheduler(
false
}
+ /**
+ * Reject group-level idioms a pipelined group cannot support, checked against the RDD graph
+ * BEFORE any stage is created -- so a rejection fails the job up front (via handleJobSubmitted's
+ * listener.jobFailed) without leaving partial scheduler state behind, exactly like the
+ * speculation check.
+ *
+ * Call only for a job that has a pipelined dependency (handleJobSubmitted gates on
+ * hasPipelined): the resource-profile check below is not keyed on a pipelined dependency, so on a
+ * regular job it would reject an ordinary RDD.withResources(...) use. Throws
+ * PIPELINED_SHUFFLE_UNSUPPORTED on violation. Enforces:
+ * - Fan-out: a pipelined producer feeding more than one consumer. 1:N is a supported model not
+ * yet built (it needs multicast to N live readers), so it is rejected for now. A
+ * PipelinedShuffleDependency's producer is `dep.rdd`; a "consumer" is any RDD that lists that
+ * dependency. More than one distinct consumer RDD for the same pipelined shuffle is fan-out.
+ * - Reliable RDD checkpoint in a group member's within-stage chain (producer OR consumer side):
+ * a reliable `checkpoint()` writes a durable, lineage-truncated snapshot, which both
+ * reintroduces cross-time reuse of a transient edge and requires a post-success recompute
+ * of the member's transient input -- for a consumer, that input is the vanished pipelined
+ * shuffle. Checked here (not at stage creation) so a reject leaves no partial stage state, and
+ * so BOTH the producer chain (rooted at pd.rdd) and each consumer chain (rooted at a consuming
+ * RDD) are covered from the whole-graph view.
+ * - A non-default resource profile on any member. The gang slot check compares one demand
+ * against one profile's capacity and measures it against the default profile, so the whole
+ * group is required to run on the default profile; any member with an explicit non-default
+ * profile is rejected. Per-profile accounting is a follow-up.
+ *
+ * (The remaining group-level case -- a regular shuffle internal to a group -- is a structural
+ * invariant that does not arise for the prefix -> pipelined-group -> suffix shapes targeted here:
+ * groups are split at regular-shuffle boundaries. The producer-side idioms -- barrier, DRA,
+ * indeterminate, checksum, push-merge -- are rejected in checkPipelinedProducerSupported at stage
+ * creation, where a producer-only throw leaves no partial state.)
+ */
+ private def checkPipelinedGroupsSupportedInRDDGraph(finalRDD: RDD[_]): Unit = {
+ // Walk the whole RDD graph once, collecting for each pipelined shuffleId the distinct consumer
+ // RDDs that read it (for the fan-out check), the producer RDDs that write it (roots of producer
+ // member stages), and every reliably-checkpointed RDD (to locate ones inside a member stage).
+ val consumersByShuffleId = new HashMap[Int, HashSet[Int]]
+ val producerRoots = new HashSet[RDD[_]] // RDDs that WRITE a pipelined shuffle
+ val reliablyCheckpointed = new HashSet[RDD[_]] // RDDs with a reliable checkpoint pending
+ var hasNonDefaultResourceProfile = false // any member RDD with a non-default RP
+ traverseRDDGraph(finalRDD) { (rdd, enqueue) =>
+ if (rdd.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]])) {
+ reliablyCheckpointed += rdd
+ }
+ // An RDD's EFFECTIVE profile is its explicit one, or the default when unset. The slot check
+ // measures capacity against the default profile (see rejectUnadmittablePipelinedGroup), so
+ // the whole group must run on the default profile; any explicit non-default profile on a
+ // member makes the group span profiles (against the default the rest use) and is rejected.
+ val rp = rdd.getResourceProfile()
+ if (rp != null && rp.id != ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID) {
+ hasNonDefaultResourceProfile = true
+ }
+ rdd.dependencies.foreach {
+ case pd: PipelinedShuffleDependency[_, _, _] =>
+ consumersByShuffleId.getOrElseUpdate(pd.shuffleId, new HashSet[Int]) += rdd.id
+ producerRoots += pd.rdd
+ enqueue(pd.rdd)
+ case dep =>
+ enqueue(dep.rdd)
+ }
+ }
+ if (consumersByShuffleId.values.exists(_.size > 1)) {
+ throw pipelinedUnsupportedError(
+ "a pipelined producer with more than one consumer (fan-out / branching)")
+ }
+ // Resource profile. The gang slot check compares one demand against one profile's capacity
+ // (maxNumConcurrentTasks is defined per profile) and measures it against the DEFAULT profile,
+ // so the whole group is required to run on the default profile and fails fast otherwise;
+ // per-profile accounting is a follow-up. Reject if ANY member carries an explicit non-default
+ // profile -- that both spans profiles (against the default the other members use) and would be
+ // admitted against the wrong (default) capacity pool. This is a real check, not just a
+ // documented assumption: nothing else enforces it.
+ if (hasNonDefaultResourceProfile) {
+ throw pipelinedUnsupportedError(
+ "a pipelined group member with a non-default resource profile (the whole group must run " +
+ "on the default profile)")
+ }
+ // Reject a reliable checkpoint anywhere in a pipelined-group MEMBER's within-stage chain.
+ // Keyed on checkpointData being ReliableRDDCheckpointData, not isCheckpointed, since the
+ // write has not happened yet. Cache / .persist() / local checkpoint are whole-partition and
+ // ephemeral and are not rejected. A reliably-checkpointed RDD `cp` is inside a member stage
+ // iff:
+ // - PRODUCER side: cp is within some producer root's own within-stage chain (walk parents from
+ // the producer root, stopping at shuffle boundaries), OR
+ // - CONSUMER side: cp's OWN within-stage chain reads a pipelined shuffle (walk parents from
+ // cp, stopping at shuffle boundaries, and check whether any stopped-at boundary is
+ // pipelined).
+ // Rooting the consumer check at each checkpointed RDD (rather than at the PSD-reading RDD) is
+ // what makes it cover a checkpoint anywhere DOWNSTREAM in the consumer stage, not just on the
+ // reading RDD itself.
+ def chainHasReliableCheckpoint(root: RDD[_]): Boolean =
+ !traverseParentRDDsWithinStage(root, (r: RDD[_]) =>
+ !r.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]]))
+ val offending =
+ // CONSUMER side: a checkpointed RDD whose own within-stage chain reads a pipelined shuffle is
+ // inside a consumer member stage (covers a checkpoint anywhere in that stage, not just on the
+ // reading RDD). PRODUCER side: a producer root's within-stage chain carries a checkpoint.
+ reliablyCheckpointed.exists(rddChainReadsPipelinedShuffle) ||
+ producerRoots.exists(chainHasReliableCheckpoint)
+ if (offending) {
+ throw pipelinedUnsupportedError(
+ "a reliable RDD checkpoint in a pipelined-group member's within-stage chain")
+ }
+ }
+
+ /** Whether `rdd`'s within-stage chain (parents, stopping at shuffle boundaries) reads through a
+ * [[PipelinedShuffleDependency]] -- i.e. `rdd` is inside a pipelined CONSUMER member stage. */
+ private def rddChainReadsPipelinedShuffle(rdd: RDD[_]): Boolean = {
+ !traverseRDDGraphUntil(rdd) { (r, enqueue) =>
+ val readsPipelined = r.dependencies.exists {
+ case _: PipelinedShuffleDependency[_, _, _] => true
+ case _: ShuffleDependency[_, _, _] => false // regular boundary: not within this stage
+ case narrowDep =>
+ enqueue(narrowDep.rdd)
+ false
+ }
+ !readsPipelined
+ }
+ }
+
/** Invoke `.partitions` on the given RDD and all of its ancestors */
private def eagerlyComputePartitionsForRddAndAncestors(rdd: RDD[_]): Unit = {
val startTime = System.nanoTime
@@ -1574,49 +1772,17 @@ private[spark] class DAGScheduler(
*/
private def rejectUnadmittablePipelinedGroup(
jobId: Int, finalRDD: RDD[_], partitions: Array[Int], listener: JobListener): Boolean = {
- // Reject up-front (before any stage is created) two group members the admission model cannot
- // support, walking the group's RDD graph once:
- // - A barrier member. A barrier stage exposes its output only after a global sync, which
- // contradicts a pipelined consumer reading the producer's output incrementally as it runs;
- // and its failure recovery resubmits the stage, which the group's atomic completion cannot
- // accommodate (a resubmitted producer would drop a co-scheduled consumer's buffered
- // completions). Reject it here rather than let it be co-scheduled.
- // - A member on a non-default resource profile. Admission below measures capacity and
- // occupancy against the default profile, but each stage derives its profile from its RDDs
- // (see createShuffleMapStage/createResultStage). A member carrying a non-default profile would be
- // admitted against the default profile's free slots yet run in a different, often smaller
- // pool -- and could then queue or deadlock there.
- var offendingBarrier = false
- var offendingRp: Option[ResourceProfile] = None
- traverseRDDGraph(finalRDD) { (rdd, enqueue) =>
- if (rdd.isBarrier()) {
- offendingBarrier = true
- }
- val rp = rdd.getResourceProfile()
- if (offendingRp.isEmpty && rp != null && rp.id != DEFAULT_RESOURCE_PROFILE_ID) {
- offendingRp = Some(rp)
- }
- rdd.dependencies.foreach(dep => enqueue(dep.rdd))
- }
- if (offendingBarrier) {
- logWarning(log"Rejecting job ${MDC(JOB_ID, jobId)}: a pipelined stage group contains a " +
- log"barrier member")
- listener.jobFailed(new SparkException(
- "A pipelined shuffle job with a barrier stage in the pipelined group is not supported: a " +
- "barrier stage exposes its output only after a global sync, which is incompatible with " +
- "a pipelined consumer reading its output incrementally."))
- return true
- }
- if (offendingRp.nonEmpty) {
- logWarning(log"Rejecting job ${MDC(JOB_ID, jobId)}: a pipelined stage group member uses a " +
- log"non-default resource profile")
- listener.jobFailed(new SparkException(
- "A pipelined shuffle job with a member on a non-default resource profile is not " +
- s"supported (resource profile id ${offendingRp.get.id}): the whole pipelined stage " +
- "group must run on the default resource profile."))
- return true
- }
if (!sc.conf.get(config.PIPELINED_GROUP_SLOT_CHECK_ENABLED)) {
+ // The only deadlock-prevention check for gang admission is off. Legitimate only when the
+ // deployment admits capacity out-of-band (e.g. a slot reservation); otherwise a pipelined
+ // group that cannot co-fit will be gang-scheduled and can deadlock. Warn once so this is
+ // never a silent state.
+ if (!warnedPipelinedSlotCheckDisabled) {
+ warnedPipelinedSlotCheckDisabled = true
+ logWarning(log"${MDC(CONFIG, config.PIPELINED_GROUP_SLOT_CHECK_ENABLED.key)}=false: " +
+ log"pipelined-group gang admission is NOT checking free slots. This is safe only if " +
+ log"capacity is reserved out-of-band; otherwise a group that cannot co-fit may deadlock.")
+ }
return false
}
val rp = sc.resourceProfileManager.defaultResourceProfile
@@ -1680,24 +1846,12 @@ private[spark] class DAGScheduler(
* member's `TaskSet.isPipelined` at submission and routing a member FetchFailed to a whole-group
* abort.
*/
- private def isPipelinedGroupMember(stage: Stage): Boolean = {
- if (isPipelinedProducer(stage)) {
- return true
- }
- // Consumer check: does this stage read through a PipelinedShuffleDependency at one of its
- // shuffle boundaries? Walk the stage's own RDD graph (descending narrow deps, stopping at every
- // shuffle boundary) -- the same edges that define the stage -- and look for a pipelined one.
- !traverseRDDGraphUntil(stage.rdd) { (rdd, enqueue) =>
- val hasPipelinedBoundary = rdd.dependencies.exists {
- case _: PipelinedShuffleDependency[_, _, _] => true
- case _: ShuffleDependency[_, _, _] => false // regular shuffle boundary: do not descend
- case narrowDep =>
- enqueue(narrowDep.rdd)
- false
- }
- !hasPipelinedBoundary // keep walking until a pipelined boundary is found
- }
- }
+ private def isPipelinedGroupMember(stage: Stage): Boolean =
+ // Producer side: it writes a pipelined shuffle. Consumer side: its within-stage chain reads
+ // one. rddChainReadsPipelinedShuffle is the single source of truth for that within-stage walk
+ // (descend narrow deps, stop at every shuffle boundary, look for a pipelined one) -- do not
+ // re-inline it; the consumer walk here and that method used to be byte-identical copies.
+ isPipelinedProducer(stage) || rddChainReadsPipelinedShuffle(stage.rdd)
/** Finds the earliest-created active job that needs the stage */
// TODO: Probably should actually find among the active jobs that need this
@@ -1875,9 +2029,19 @@ private[spark] class DAGScheduler(
if (hasPipelined && rejectUnadmittablePipelinedGroup(jobId, finalRDD, partitions, listener)) {
return
}
-
var finalStage: ResultStage = null
try {
+ // Reject group-level unsupported pipelined idioms (e.g. fan-out, a non-default resource
+ // profile, a reliable checkpoint in a member stage) from the RDD graph, up front -- before
+ // any stage is created, so a rejection leaves no partial scheduler state. Gated on
+ // hasPipelined: every idiom this checks concerns a pipelined group, so it must not run for a
+ // job with no pipelined dependency (the resource-profile check in particular is not keyed on
+ // a pipelined dependency and would otherwise reject an ordinary job that merely uses a
+ // non-default profile via RDD.withResources). Inside this try so any incidental exception
+ // from the graph walk is handled by the same listener.jobFailed path as stage creation.
+ if (hasPipelined) {
+ checkPipelinedGroupsSupportedInRDDGraph(finalRDD)
+ }
// New stage creation may throw an exception if, for example, jobs are run on a
// HadoopRDD whose underlying HDFS files have been deleted.
finalStage = createResultStage(finalRDD, func, partitions, jobId, callSite)
@@ -1910,6 +2074,15 @@ private[spark] class DAGScheduler(
return
}
+ case e: PipelinedShuffleUnsupportedException =>
+ // An up-front idiom rejection (checkPipelinedGroupsSupportedInRDDGraph / a producer-side
+ // check in createShuffleMapStage), 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)
@@ -1918,7 +2091,9 @@ private[spark] class DAGScheduler(
// Job submitted, clear internal data.
barrierJobIdToNumTasksCheckFailures.remove(jobId)
- val job = new ActiveJob(jobId, finalStage, callSite, listener, artifacts, properties)
+ // Pass hasPipelined (computed above) into the job; see ActiveJob.hasPipelinedDependency.
+ val job = new ActiveJob(jobId, finalStage, callSite, listener, artifacts, properties,
+ hasPipelinedDependency = hasPipelined)
clearCacheLocs()
logInfo(
log"Got job ${MDC(JOB_ID, job.jobId)} (${MDC(CALL_SITE_SHORT_FORM, callSite.shortForm)}) " +
@@ -2077,7 +2252,9 @@ private[spark] class DAGScheduler(
// rejectUnadmittablePipelinedGroup) before any member was submitted, so the group is
// known to fit; just co-schedule this consumer with its running producer(s). No slot
// check here -- that would re-measure capacity against a mid-flight snapshot and is
- // unnecessary once admission is decided up front (gang admission).
+ // unnecessary once admission is decided up front (gang admission). Group-level
+ // idiom rejection (fan-out, internal regular shuffle) already happened at job
+ // submission (checkPipelinedGroupsSupportedInRDDGraph + the all-pipelined check).
logInfo(log"Submitting ${MDC(STAGE, stage)} concurrently with its running " +
log"pipelined producer(s) ${MDC(MISSING_PARENT_STAGES, pipelinedMissing)}")
// Record that this stage is co-scheduled with still-running pipelined producers,
@@ -2538,9 +2715,13 @@ private[spark] class DAGScheduler(
case _: ResultStage => None
}
+ // Only a job that uses a pipelined shuffle can have a pipelined-group member; gate the
+ // group-membership graph walk on that cheap per-job flag so a regular job pays nothing here.
+ val isPipelined = jobIdToActiveJob.get(jobId).exists(_.hasPipelinedDependency) &&
+ isPipelinedGroupMember(stage)
taskScheduler.submitTasks(new TaskSet(
tasks.toArray, stage.id, stage.latestInfo.attemptNumber(), jobId, properties,
- stage.resourceProfileId, shuffleId))
+ stage.resourceProfileId, shuffleId, isPipelined = isPipelined))
} else {
// Because we posted SparkListenerStageSubmitted earlier, we should mark
// the stage as completed here in case there are no tasks to run
@@ -3081,7 +3262,29 @@ private[spark] class DAGScheduler(
case smt: ShuffleMapTask =>
val shuffleStage = stage.asInstanceOf[ShuffleMapStage]
- if (!ignoreOldTaskAttempts) {
+ if (shuffleStage.isPipelined) {
+ // A pipelined shuffle's completed partitions are tracked locally and monotonically on
+ // the stage, not in the MapOutputTracker (the reader finds the producer via the
+ // streaming transport, not the tracker). See ShuffleMapStage's
+ // pipelinedCompletedPartitions scaladoc for why -- the crux of avoiding the
+ // streaming-writer resubmit hang. Checksum-mismatch detection does not apply (a
+ // dependency never enables checksum retry -- see PipelinedShuffleDependency).
+ //
+ // Record the partition and decrement pendingPartitions UNCONDITIONALLY -- outside the
+ // `!ignoreOldTaskAttempts` and bogus-epoch guards that gate a regular shuffle. Both
+ // guards exist only to avoid trusting a MapOutputTracker registration that a later
+ // rollback (ignoreOldTaskAttempts, from an indeterminate/rolled-back stage) or an
+ // executor-loss strip (bogus epoch) could invalidate. A pipelined stage never
+ // registers there, and its completed set is monotonic and never rolled back (a
+ // transient shuffle cannot be recomputed; any real group failure aborts the whole
+ // group). Skipping the record for an "old" or "bogus" straggler would be actively
+ // harmful: with pendingPartitions decremented but the partition unrecorded, a dropped
+ // last partition leaves the stage "done but not available" -> processShuffleMapStage-
+ // Completion resubmits the transient producer, reopening the streaming-writer hang.
+ // An already-successful straggler is not a failure, so recording it is correct.
+ shuffleStage.pendingPartitions -= task.partitionId
+ shuffleStage.addPipelinedCompletedPartition(smt.partitionId)
+ } else if (!ignoreOldTaskAttempts) {
shuffleStage.pendingPartitions -= task.partitionId
val status = event.result.asInstanceOf[MapStatus]
val execId = status.location.executorId
@@ -3149,6 +3352,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)
@@ -3265,39 +3495,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 =>
@@ -3837,6 +4035,52 @@ private[spark] class DAGScheduler(
maybeEpoch = None)
}
+ /**
+ * On a FetchFailed, unregister the shuffle outputs of the executor (or its whole host) whose
+ * fetch failed, treating the FetchFailed as authoritative evidence that its shuffle data is gone.
+ * Extracted from the base FetchFailed handler so the pipelined-group-abort branch can also run it:
+ * aborting the group fails only this job's stages, but a dead executor's REGULAR outputs must
+ * still be cleaned up for other/concurrent jobs (with an external shuffle service, an ExecutorLost
+ * does not clean them, so FetchFailed is the only proactive channel). No-op when `bmAddress` is
+ * null. Safe for a pipelined shuffle: it registers no map outputs in the tracker, so this can only
+ * strip regular/durable outputs.
+ */
+ private def unregisterOutputsOnFetchFailedExecutor(
+ bmAddress: BlockManagerId, task: Task[_]): Unit = {
+ // TODO: mark the executor as failed only if there were lots of fetch failures on it
+ if (bmAddress != null) {
+ val externalShuffleServiceEnabled = env.blockManager.externalShuffleServiceEnabled
+ val isHostDecommissioned = taskScheduler
+ .getExecutorDecommissionState(bmAddress.executorId)
+ .exists(_.workerHost.isDefined)
+
+ // Shuffle output of all executors on host `bmAddress.host` may be lost if:
+ // - External shuffle service is enabled, so we assume that all shuffle data on node is bad.
+ // - Host is decommissioned, thus all executors on that host will die.
+ val shuffleOutputOfEntireHostLost = externalShuffleServiceEnabled || isHostDecommissioned
+ val hostToUnregisterOutputs = if (shuffleOutputOfEntireHostLost
+ && unRegisterOutputOnHostOnFetchFailure) {
+ Some(bmAddress.host)
+ } else {
+ // Unregister shuffle data just for one executor (we don't have any
+ // reason to believe shuffle data has been lost for the entire host).
+ None
+ }
+ removeExecutorAndUnregisterOutputs(
+ execId = bmAddress.executorId,
+ fileLost = true,
+ hostToUnregisterOutputs = hostToUnregisterOutputs,
+ maybeEpoch = Some(task.epoch),
+ // shuffleFileLostEpoch is ignored when a host is decommissioned because some
+ // decommissioned executors on that host might have been removed before this fetch
+ // failure and might have bumped up the shuffleFileLostEpoch. We ignore that, and
+ // proceed with unconditional removal of shuffle outputs from all executors on that
+ // host, including from those that we still haven't confirmed as lost due to heartbeat
+ // delays.
+ ignoreShuffleFileLostEpoch = isHostDecommissioned)
+ }
+ }
+
/**
* Handles removing an executor from the BlockManagerMaster as well as unregistering shuffle
* outputs for the executor or optionally its host.
@@ -4449,6 +4693,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/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 3a5f60ae1935..4fbe197cb838 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
@@ -70,6 +70,15 @@ private[spark] class TaskSetManager(
val ser = env.closureSerializer.newInstance()
val tasks = taskSet.tasks
+
+ // A pipelined-group member reads/writes a transient shuffle that cannot be re-read in isolation,
+ // so a single task failure must fail the whole group (which fails the job, triggering a rerun)
+ // rather than be retried per-task. For such a task set we cap attempts at 1 and count every
+ // failure, including reasons that are normally not counted (e.g. executor loss). This is the
+ // native equivalent of the prototype's per-batch "any failure counts" behavior, keyed on the
+ // pipelined dependency (via TaskSet.isPipelined) rather than a job property.
+ private val effectiveMaxTaskFailures = if (taskSet.isPipelined) 1 else maxTaskFailures
+
private val isShuffleMapTasks = tasks(0).isInstanceOf[ShuffleMapTask]
// shuffleId is only available when isShuffleMapTasks=true
private val shuffleId = taskSet.shuffleId
@@ -1106,16 +1115,34 @@ 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: TaskKilled (a deliberate kill) and TaskCommitDenied. These are not
+ // normally even reachable for a pipelined group -- speculation is rejected up front for a
+ // pipelined job, and a group never does an in-place stage/task retry (a member failure aborts
+ // the whole group; the caller reruns as a fresh job), so the usual sources of these reasons do
+ // not arise. The exclusion is kept as a defensive guard: if one ever did reach a live member
+ // (e.g. a manual killTaskAttempt), it must be treated as benign rather than spuriously aborting
+ // the group. So for a pipelined set we count everything EXCEPT those two benign reasons.
+ // Non-pipelined sets are unchanged.
+ val forceCountForPipelined = taskSet.isPipelined && (reason match {
+ case _: TaskKilled | _: TaskCommitDenied => false
+ case _ => true
+ })
+ val countTowardsTaskFailures = reason.countTowardsTaskFailures || forceCountForPipelined
+ if (!isZombie && countTowardsTaskFailures) {
assert (null != failureReason)
taskSetExcludelistHelperOpt.foreach(_.updateExcludedForFailedTask(
info.host, info.executorId, index, failureReasonString))
numFailures(index) += 1
- if (numFailures(index) >= maxTaskFailures) {
+ if (numFailures(index) >= effectiveMaxTaskFailures) {
logError(log"Task ${MDC(TASK_INDEX, index)} in stage " + taskSet.logId +
- log" failed ${MDC(MAX_ATTEMPTS, maxTaskFailures)} times; aborting job")
+ log" failed ${MDC(MAX_ATTEMPTS, effectiveMaxTaskFailures)} times; aborting job")
abort("Task %d in stage %s failed %d times, most recent failure: %s\nDriver stacktrace:"
- .format(index, taskSet.id, maxTaskFailures, failureReasonString), failureException)
+ .format(index, taskSet.id, effectiveMaxTaskFailures, failureReasonString),
+ failureException)
return
}
}
@@ -1191,7 +1218,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 f43354c80865..05ba990e2d63 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala
@@ -41,7 +41,7 @@ import org.apache.spark.executor.ExecutorMetrics
import org.apache.spark.internal.config
import org.apache.spark.internal.config.{LEGACY_ABORT_STAGE_AFTER_KILL_TASKS, Tests}
import org.apache.spark.network.shuffle.ExternalBlockStoreClient
-import org.apache.spark.rdd.{DeterministicLevel, RDD}
+import org.apache.spark.rdd.{DeterministicLevel, RDD, ReliableRDDCheckpointData}
import org.apache.spark.resource.{ExecutorResourceRequests, ResourceProfile, ResourceProfileBuilder, TaskResourceProfile, TaskResourceRequests}
import org.apache.spark.resource.ResourceUtils.{FPGA, GPU}
import org.apache.spark.rpc.RpcTimeoutException
@@ -6432,6 +6432,31 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti
assertDataStructuresEmpty()
}
+ test("pipelined shuffle: a regular-shuffle prefix feeding a pipelined producer is rejected") {
+ // Also mixed: a regular shuffle in the PREFIX that feeds a pipelined producer
+ // (regularRoot --regular--> producer(pipelined) --pipelined--> consumer). Rather than treat
+ // the regular edge as an ordinary external input to the group and support a mid-DAG regular
+ // prefix, a pipelined job rejects ANY regular shuffle up front (the supported shape is
+ // scan-of-files --pipelined--> stateful, with no upstream shuffle).
+ val regularRoot = new MyRDD(sc, 2, Nil)
+ val regularDep = new ShuffleDependency(regularRoot, new HashPartitioner(2))
+ val producerRdd = new MyRDD(sc, 2, List(regularDep), tracker = mapOutputTracker)
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+ val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ val failure = new java.util.concurrent.atomic.AtomicReference[Exception]()
+ val failListener = new JobListener {
+ override def taskSucceeded(index: Int, result: Any): Unit = {}
+ override def jobFailed(exception: Exception): Unit = failure.set(exception)
+ }
+ submit(consumerRdd, Array(0, 1), listener = failListener)
+
+ assert(failure.get() != null, "a pipelined job with a regular-shuffle prefix must fail")
+ assert(failure.get().getMessage.contains("all-regular or all-pipelined"),
+ s"expected a mixed-job rejection, got: ${failure.get().getMessage}")
+ assert(taskSets.isEmpty, "no stage should be submitted for a rejected mixed job")
+ assertDataStructuresEmpty()
+ }
+
test("pipelined shuffle: deep chain A->B->C is submitted fully concurrently") {
// A --pipelined--> B --pipelined--> C : all three co-scheduled (each edge is non-sequencing).
val rddA = new MyRDD(sc, 2, Nil)
@@ -6459,6 +6484,35 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti
assertDataStructuresEmpty()
}
+ test("pipelined shuffle: an explicit job cancellation cleans up a buffered consumer deferral") {
+ // A buffered deferral is the only mutable state this feature adds; it must never outlive its
+ // job. Job cancellation goes through failJobAndIndependentStages ->
+ // cleanupStateForJobAndIndependentStages, which drops the entry (as a consumer key) and removes
+ // the stage from every other consumer's pending-producer set. Verify a consumer whose
+ // completion is buffered while its producer runs leaves NO deferral behind when the job is
+ // cancelled, and no buffered success is later applied as a result.
+ val producerRdd = new MyRDD(sc, 2, Nil)
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+ val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ val jobId = submit(consumerRdd, Array(0, 1))
+ val consumerTaskSet = taskSets.find { ts =>
+ scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd
+ }.get
+
+ // Consumer finishes ahead of its producer -> its completion is buffered (a deferral exists).
+ complete(consumerTaskSet, Seq((Success, 42), (Success, 43)))
+ assert(scheduler.dependentStageMap.nonEmpty, "a deferral must exist while the producer runs")
+ assert(results.isEmpty, "the consumer result must be deferred, not applied yet")
+
+ // Cancel the job. The deferral must be torn down with everything else -- no stale entry, and
+ // the buffered success must not be replayed as a result.
+ cancel(jobId)
+ assert(scheduler.dependentStageMap.isEmpty,
+ "job cancellation must clean up the buffered consumer deferral (no state outlives the job)")
+ assert(results.isEmpty, "a cancelled job's buffered consumer success must not be applied")
+ assertDataStructuresEmpty()
+ }
+
test("regular shuffle job with speculation enabled is NOT rejected (rejection path is inert)") {
// The speculation fail-fast must apply only to jobs with a pipelined dependency; a plain
// regular-shuffle job with speculation on runs normally.
@@ -6532,65 +6586,6 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti
}
}
- test("pipelined shuffle: a member on a non-default resource profile is rejected up front") {
- // Admission measures capacity/occupancy against the DEFAULT resource profile, but a stage
- // derives its profile from its RDDs. A pipelined group member on a custom profile would be
- // admitted against the default pool's free slots yet run in a different (often smaller) pool
- // and could deadlock there. Reject such a job before any stage is created.
- // Ensure the default profile exists (id 0) before building a custom one, so the custom profile
- // gets a distinct, non-default id regardless of test-execution order (profile ids come from a
- // process-wide counter, and the default profile occupies id 0).
- sc.resourceProfileManager.defaultResourceProfile
- val ereqs = new ExecutorResourceRequests().cores(4)
- val treqs = new TaskResourceRequests().cpus(2)
- val customRp = new ResourceProfileBuilder().require(ereqs).require(treqs).build()
- val producerRdd = new MyRDD(sc, 2, Nil)
- val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
- val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
- .withResources(customRp)
- // Sanity: the consumer really carries a non-default profile (otherwise the test is vacuous).
- assert(consumerRdd.getResourceProfile() != null &&
- consumerRdd.getResourceProfile().id != ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID,
- s"test setup: expected a non-default profile, got id " +
- s"${Option(consumerRdd.getResourceProfile()).map(_.id)}")
- val failure = new java.util.concurrent.atomic.AtomicReference[Exception]()
- val failListener = new JobListener {
- override def taskSucceeded(index: Int, result: Any): Unit = {}
- override def jobFailed(exception: Exception): Unit = failure.set(exception)
- }
- submit(consumerRdd, Array(0, 1), listener = failListener)
- assert(failure.get() != null,
- "a pipelined job with a non-default-resource-profile member should be rejected")
- assert(failure.get().getMessage.contains("non-default resource profile"))
- assert(taskSets.isEmpty, "no stage should be created for a rejected job")
- assertDataStructuresEmpty()
- }
-
- test("pipelined shuffle: a barrier group member is rejected up front") {
- // A barrier stage exposes its output only after a global sync, which is incompatible with a
- // pipelined consumer reading the producer's output incrementally. It is also a recovery hazard:
- // a barrier task failure fails the stage and RESUBMITS it (markStageAsFinished without
- // willRetry), and if such a producer were co-scheduled, the resubmit would drop a fan-in
- // consumer's buffered completions and the job could never complete. Reject any job whose
- // pipelined group contains a barrier member before any stage is created, so that path is
- // unreachable.
- val producerRdd = new MyRDD(sc, 2, Nil).barrier().mapPartitions(iter => iter)
- assert(producerRdd.isBarrier(), "test setup: the producer must be a barrier RDD")
- val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
- val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
- val failure = new java.util.concurrent.atomic.AtomicReference[Exception]()
- val failListener = new JobListener {
- override def taskSucceeded(index: Int, result: Any): Unit = {}
- override def jobFailed(exception: Exception): Unit = failure.set(exception)
- }
- submit(consumerRdd, Array(0, 1), listener = failListener)
- assert(failure.get() != null,
- "a pipelined job with a barrier group member should be rejected")
- assert(failure.get().getMessage.contains("barrier"))
- assert(taskSets.isEmpty, "no stage should be created for a rejected job")
- assertDataStructuresEmpty()
- }
-
test("regular shuffle job with dynamic allocation enabled is NOT rejected (path is inert)") {
// The dynamic-allocation fail-fast must apply only to jobs with a pipelined dependency; a plain
// regular-shuffle job with dynamic allocation on runs normally.
@@ -6832,13 +6827,19 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti
}
}
- test("pipelined shuffle: a diamond reusing one producer counts its tasks once for admission") {
+ test("pipelined shuffle: admission counts a reused producer once (diamond rejected for " +
+ "fan-out, not for slots)") {
// Fan-out/diamond: ONE PipelinedShuffleDependency (one producer, one shuffle id) is read by TWO
- // consumers that are then narrow-joined into the result. Execution creates ONE producer stage
- // (getOrCreateShuffleMapStage keys on shuffle id), so the group's real concurrent demand is
- // producer(2) + result(2) = 4 -- NOT 6. Counting the producer once per consumer EDGE would make
- // demand 6 and wrongly reject this group at capacity 4. Pinning capacity to exactly 4 admits
- // the correctly-deduped group and would fail if the producer were double-counted.
+ // consumers that are then narrow-joined into the result. Fan-out is unsupported, so the group
+ // is ultimately rejected -- but WHICH rejection it gets proves the admission demand is deduped.
+ // The up-front slot admission (rejectUnadmittablePipelinedGroup) runs BEFORE the fan-out idiom
+ // check (checkPipelinedGroupsSupportedInRDDGraph). Execution creates ONE producer stage
+ // (getOrCreateShuffleMapStage keys on shuffle id), so the real concurrent demand is
+ // producer(2) + result(2) = 4. Counting the producer once per consumer EDGE would inflate it
+ // to 6. Pinning capacity to exactly 4: the deduped group PASSES the slot check and is then
+ // rejected for fan-out (PIPELINED_SHUFFLE_UNSUPPORTED); a double-counted group (6 > 4) would
+ // instead be rejected earlier for CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT. So a dedup regression
+ // flips the error class -- which this test detects.
val producerRdd = new MyRDD(sc, 2, Nil)
val myScheduler = scheduler.asInstanceOf[MyDAGScheduler]
myScheduler.maxConcurrentTasksForTest = 4
@@ -6852,13 +6853,13 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti
val resultRdd = new MyRDD(
sc, 2, List(new OneToOneDependency(consumer1), new OneToOneDependency(consumer2)),
tracker = mapOutputTracker)
- submit(resultRdd, Array(0, 1))
- assert(taskSets.size === 2,
- s"diamond group is producer + result (both consumers); got ${taskSets.size} task sets")
- // Drain to completion: producer stage first, then the result stage.
- completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2)
- complete(taskSets(1), Seq((Success, 42), (Success, 43)))
- assert(results === Map(0 -> 42, 1 -> 43))
+ val failure = submitAndCaptureFailure(resultRdd, Array(0, 1))
+ // Rejected for fan-out (demand fit within capacity 4), NOT for slots -- proving the reused
+ // producer was counted once.
+ assertPipelinedUnsupported(failure, "more than one consumer")
+ assert(!failure.getMessage.contains("CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT"),
+ s"a reused producer must be counted once (fit capacity 4), got a slot rejection: " +
+ s"${failure.getMessage}")
assertDataStructuresEmpty()
} finally {
myScheduler.maxConcurrentTasksForTest = 1000
@@ -6866,10 +6867,13 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti
}
}
- test("pipelined shuffle: a wide fan-out reusing one producer counts its tasks once") {
+ test("pipelined shuffle: admission counts a reused producer once across a wide fan-out " +
+ "(rejected for fan-out, not for slots)") {
// Wider fan-out: ONE producer feeds THREE consumers, all narrow-joined into the result. Real
- // demand is still producer(2) + result(2) = 4; a per-edge count would be 2 + 3*2 = 8. Capacity
- // 4 admits the deduped group and would reject the double-counted one.
+ // demand is still producer(2) + result(2) = 4; a per-edge count would be 2 + 3*2 = 8. At
+ // capacity 4 the deduped group passes the slot check and is rejected for fan-out; a
+ // double-counted group (8 > 4) would be rejected first for insufficient slots. See the diamond
+ // test above for the ordering rationale.
val producerRdd = new MyRDD(sc, 2, Nil)
val myScheduler = scheduler.asInstanceOf[MyDAGScheduler]
myScheduler.maxConcurrentTasksForTest = 4
@@ -6880,12 +6884,11 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti
new OneToOneDependency(new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker))
}.toList
val resultRdd = new MyRDD(sc, 2, consumers, tracker = mapOutputTracker)
- submit(resultRdd, Array(0, 1))
- assert(taskSets.size === 2,
- s"wide fan-out group is producer + result; got ${taskSets.size} task sets")
- completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2)
- complete(taskSets(1), Seq((Success, 42), (Success, 43)))
- assert(results === Map(0 -> 42, 1 -> 43))
+ val failure = submitAndCaptureFailure(resultRdd, Array(0, 1))
+ assertPipelinedUnsupported(failure, "more than one consumer")
+ assert(!failure.getMessage.contains("CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT"),
+ s"a reused producer must be counted once (fit capacity 4), got a slot rejection: " +
+ s"${failure.getMessage}")
assertDataStructuresEmpty()
} finally {
myScheduler.maxConcurrentTasksForTest = 1000
@@ -7075,6 +7078,45 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti
s"a dropped fan-in consumer must not replay when a surviving producer succeeds; got $results")
}
+ test("pipelined shuffle: a multi-producer consumer's buffered successes are dropped when one " +
+ "producer fails (no stale replay via a sibling)") {
+ // Group-atomic drop across MULTIPLE producers. Consumer C reads TWO pipelined producers P1 and
+ // P2 (a join; purely all-pipelined, no regular shuffle). C finishes early -> its successes are
+ // buffered against BOTH producers. If P1 then fails, the whole group must be torn down and C's
+ // successes DROPPED -- they depended on P1's (now-invalid) output. Note the release path
+ // evaluates producerFailed per finishing producer, so a naive impl could remove P1
+ // (parents still has P2, no drop), then later see P2 "succeed" and REPLAY. This asserts the
+ // shipped stack drops instead: a member failure aborts the whole group, which tears C down.
+ val rddP1 = new MyRDD(sc, 2, Nil)
+ val psdP1 = new PipelinedShuffleDependency(rddP1, new HashPartitioner(2))
+ val rddP2 = new MyRDD(sc, 2, Nil)
+ val psdP2 = new PipelinedShuffleDependency(rddP2, new HashPartitioner(2))
+ val consumerRdd = new MyRDD(sc, 2, List(psdP1, psdP2), tracker = mapOutputTracker)
+ val failure = new java.util.concurrent.atomic.AtomicReference[Exception]()
+ val failListener = new JobListener {
+ override def taskSucceeded(index: Int, result: Any): Unit = results.put(index, result)
+ override def jobFailed(exception: Exception): Unit = failure.set(exception)
+ }
+ submit(consumerRdd, Array(0, 1), listener = failListener)
+ // P1, P2, and C are all co-scheduled (purely all-pipelined join, admitted up front).
+ assert(taskSets.size === 3, s"expected P1, P2, consumer co-scheduled, got ${taskSets.size}")
+ val tsP1 = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddP1).get
+ val tsC = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd).get
+
+ // Consumer finishes early -> buffered against both P1 and P2.
+ complete(tsC, Seq((Success, 42), (Success, 43)))
+ assert(results.isEmpty, "consumer completions should be buffered while producers run")
+ assert(scheduler.dependentStageMap.keys.exists(_.rdd eq consumerRdd),
+ "consumer should be deferred against its two producers")
+
+ // P1 fails. The group is torn down; C's buffered successes must be DROPPED, never replayed --
+ // even though P2 has not (and now will not) complete.
+ failed(tsP1, "producer P1 blew up")
+ assert(failure.get() != null, "the job must fail when a pipelined producer fails")
+ assert(results.isEmpty,
+ "a multi-producer consumer's buffered successes must be dropped when any producer fails")
+ assertDataStructuresEmpty()
+ }
test("pipelined shuffle: a deferred consumer task fires its TaskEnd exactly once (at replay)") {
// A deferred CompletionEvent must have its side effects (task-end listener event, accumulator
@@ -7219,6 +7261,772 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti
}
}
+ test("pipelined shuffle: both producer and consumer task sets are marked isPipelined") {
+ // The TaskSet.isPipelined flag drives group-atomic failure in the task scheduler; verify the
+ // DAGScheduler sets it for both members of a pipelined group (producer and consumer).
+ val producerRdd = new MyRDD(sc, 2, Nil)
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+ val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ submit(consumerRdd, Array(0, 1))
+ assert(taskSets.size === 2)
+ val producerTs =
+ taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq producerRdd).get
+ val consumerTs =
+ taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd).get
+ assert(producerTs.isPipelined, "the pipelined producer's task set must be marked isPipelined")
+ assert(consumerTs.isPipelined, "the pipelined consumer's task set must be marked isPipelined")
+
+ completeShuffleMapStageSuccessfully(producerTs.stageId, 0, 2)
+ complete(consumerTs, Seq((Success, 42), (Success, 43)))
+ assert(results === Map(0 -> 42, 1 -> 43))
+ assertDataStructuresEmpty()
+ }
+
+ test("regular shuffle: task sets are NOT marked isPipelined (inertness)") {
+ // A regular producer/consumer must not be marked isPipelined.
+ val producerRdd = new MyRDD(sc, 2, Nil)
+ val regularDep = new ShuffleDependency(producerRdd, new HashPartitioner(2))
+ val consumerRdd = new MyRDD(sc, 2, List(regularDep), tracker = mapOutputTracker)
+ submit(consumerRdd, Array(0, 1))
+ assert(taskSets.head.isPipelined === false, "a regular producer must not be marked isPipelined")
+ completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2)
+ assert(taskSets(1).isPipelined === false, "a regular consumer must not be marked isPipelined")
+ complete(taskSets(1), Seq((Success, 42), (Success, 43)))
+ assert(results === Map(0 -> 42, 1 -> 43))
+ assertDataStructuresEmpty()
+ }
+
+ // ==========================================================================================
+ // Cross-job / cross-time reuse prevention at both layers: a consumed pipelined
+ // producer whose executor is lost must not be resubmitted
+ // ==========================================================================================
+
+ test("pipelined shuffle: producer availability is tracked on the stage, not the" +
+ "MapOutputTracker") {
+ // A pipelined producer's completed partitions are tracked on ShuffleMapStage (monotonic), not
+ // registered as durable outputs in the MapOutputTracker. Verify: after the producer's map tasks
+ // succeed, the stage is available WITHOUT the shuffle having map outputs in the tracker.
+ val producerRdd = new MyRDD(sc, 2, Nil)
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+ val shuffleId = pipelinedDep.shuffleId
+ val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ submit(consumerRdd, Array(0, 1))
+ val producerStageId = taskSets.head.stageId
+ val producerStage =
+ scheduler.stageIdToStage(producerStageId).asInstanceOf[ShuffleMapStage]
+
+ // Complete the producer's two map tasks.
+ complete(taskSets.head, Seq(
+ (Success, makeMapStatus("hostA", 2)),
+ (Success, makeMapStatus("hostB", 2))))
+
+ // The stage is available via its local completed-partition set...
+ assert(producerStage.isAvailable,
+ "pipelined producer should be available after its tasks finish")
+ assert(producerStage.isPipelined)
+ // ...but the pipelined shuffle has NO map outputs registered in the MapOutputTracker.
+ assert(mapOutputTracker.getNumAvailableOutputs(shuffleId) === 0,
+ "a pipelined shuffle must not register durable map outputs in the MapOutputTracker")
+
+ complete(taskSets(1), Seq((Success, 42), (Success, 43)))
+ assert(results === Map(0 -> 42, 1 -> 43))
+ assertDataStructuresEmpty()
+ }
+
+ test("pipelined shuffle: numAvailableOutputs/findMissingPartitions reflect the exact " +
+ "incomplete partition set") {
+ // Pin the availability CONTRACT for a partially-complete pipelined producer, not just the
+ // all-done / none-done endpoints: numAvailableOutputs must be the count of completed partitions
+ // and findMissingPartitions must return the exact ids still missing (identity, not just size).
+ // Without this, a plausible refactor (revert the findMissingPartitions isPipelined branch to
+ // the
+ // tracker, use a bare counter, drop/take by size, or invert the filter) would silently break
+ // the
+ // contract yet pass every full-completion test.
+ val producerRdd = new MyRDD(sc, 2, Nil)
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+ val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ submit(consumerRdd, Array(0, 1))
+ val producerTs = taskSets.head
+ val producerStage =
+ scheduler.stageIdToStage(producerTs.stageId).asInstanceOf[ShuffleMapStage]
+
+ // Complete ONLY partition 1, leaving partition 0 missing.
+ runEvent(makeCompletionEvent(producerTs.tasks(1), Success, makeMapStatus("hostB", 2)))
+ assert(producerStage.numAvailableOutputs === 1)
+ assert(!producerStage.isAvailable)
+ assert(producerStage.findMissingPartitions() === Seq(0),
+ "findMissingPartitions must name the exact missing partition, not just a right-sized set")
+
+ // Complete partition 0; now nothing is missing.
+ runEvent(makeCompletionEvent(producerTs.tasks(0), Success, makeMapStatus("hostA", 2)))
+ assert(producerStage.numAvailableOutputs === 2)
+ assert(producerStage.isAvailable)
+ assert(producerStage.findMissingPartitions() === Seq.empty)
+
+ // Drain the consumer cleanly.
+ val consumerTs = taskSets.find { ts =>
+ scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd
+ }.get
+ complete(consumerTs, Seq((Success, 42), (Success, 43)))
+ assert(results === Map(0 -> 42, 1 -> 43))
+ assertDataStructuresEmpty()
+ }
+
+ test("pipelined shuffle: a post-executor-loss straggler success must not resubmit the producer " +
+ "(bogus-epoch race)") {
+ // A pipelined producer's map task can succeed on an executor whose loss is already recorded
+ // (its StatusUpdate raced the executor-loss event). That completion hits the "possibly bogus
+ // epoch" branch, which for a regular shuffle simply ignores it (a healthy reattempt will
+ // re-register the output). But the same branch also runs `pendingPartitions -= partitionId`
+ // first, so if it is the last pending partition and we do NOT record it in the pipelined
+ // completed set, the stage looks "done but not available" and processShuffleMapStageCompletion
+ // resubmits the transient producer -- the exact streaming-writer hang. A pipelined stage must
+ // record the partition as completed even on the bogus-epoch path (its output is monotonic and
+ // MapOutputTracker's executor-loss stripping does not apply to it).
+ val producerRdd = new MyRDD(sc, 2, Nil)
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+ val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ submit(consumerRdd, Array(0, 1))
+ val producerTs = taskSets.head
+ val producerStage =
+ scheduler.stageIdToStage(producerTs.stageId).asInstanceOf[ShuffleMapStage]
+
+ // Partition 0 completes normally on executor "hostA-exec" (recorded). Note makeMapStatus(host)
+ // builds executorId = host + "-exec", so pass host "hostA" to get executorId "hostA-exec".
+ runEvent(makeCompletionEvent(producerTs.tasks(0), Success, makeMapStatus("hostA", 2)))
+ assert(producerStage.numAvailableOutputs === 1)
+ val taskSetsBefore = taskSets.size
+
+ // Executor "hostA-exec" is lost: this records executorFailureEpoch("hostA-exec") = current
+ // epoch.
+ runEvent(ExecutorLost("hostA-exec", ExecutorExited(-100, false, "Container marked as failed")))
+
+ // Now a delayed straggler Success for partition 1 arrives FROM the lost executor "hostA-exec".
+ // Its task epoch is the default (-1) <= the recorded failure epoch, so it is treated as
+ // possibly-bogus. It must still be recorded for a pipelined stage, so the producer becomes
+ // available and is NOT resubmitted.
+ runEvent(makeCompletionEvent(producerTs.tasks(1), Success, makeMapStatus("hostA", 2)))
+
+ assert(producerStage.isAvailable,
+ "a pipelined producer must be available after all partitions succeed, even via a " +
+ "bogus-epoch " +
+ "straggler")
+ assert(producerStage.findMissingPartitions() === Seq.empty)
+ assert(taskSets.size === taskSetsBefore,
+ "the pipelined producer must NOT be resubmitted by a post-loss straggler success")
+
+ // Consumer drains normally; no hang.
+ val consumerTs = taskSets.find { ts =>
+ scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd
+ }.get
+ complete(consumerTs, Seq((Success, 42), (Success, 43)))
+ assert(results === Map(0 -> 42, 1 -> 43))
+ assertDataStructuresEmpty()
+ }
+
+ test("pipelined shuffle: a producer success with ignoreOldTaskAttempts set is still recorded " +
+ "(no resubmit)") {
+ // Sibling of the bogus-epoch case, for the OTHER guard that gates a regular shuffle's
+ // recording:
+ // ignoreOldTaskAttempts (set when a stage is rolled back, e.g. as a succeeding stage of an
+ // indeterminate ancestor). A pipelined producer's completed set is monotonic and never rolled
+ // back, so its success must be recorded even when ignoreOldTaskAttempts is true -- otherwise
+ // the
+ // last partition is dropped (pendingPartitions decremented, not recorded) -> "done but not
+ // available" -> processShuffleMapStageCompletion resubmits the transient producer.
+ val producerRdd = new MyRDD(sc, 2, Nil)
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+ val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ submit(consumerRdd, Array(0, 1))
+ val producerTs = taskSets.head
+ val producerStage =
+ scheduler.stageIdToStage(producerTs.stageId).asInstanceOf[ShuffleMapStage]
+
+ runEvent(makeCompletionEvent(producerTs.tasks(0), Success, makeMapStatus("hostA", 2)))
+ assert(producerStage.numAvailableOutputs === 1)
+ val taskSetsBefore = taskSets.size
+
+ // Force ignoreOldTaskAttempts=true for the next completion: maxAttemptIdToIgnore >= the task's
+ // stageAttemptId (0). For a regular stage this would drop the completion; a pipelined stage
+ // must
+ // still record it.
+ producerStage.maxAttemptIdToIgnore = Some(0)
+ runEvent(makeCompletionEvent(producerTs.tasks(1), Success, makeMapStatus("hostB", 2)))
+
+ assert(producerStage.isAvailable,
+ "a pipelined producer success must be recorded even when ignoreOldTaskAttempts is set")
+ assert(producerStage.findMissingPartitions() === Seq.empty)
+ assert(taskSets.size === taskSetsBefore,
+ "the pipelined producer must NOT be resubmitted when a success arrives under " +
+ "ignoreOldTaskAttempts")
+
+ val consumerTs = taskSets.find { ts =>
+ scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd
+ }.get
+ complete(consumerTs, Seq((Success, 42), (Success, 43)))
+ assert(results === Map(0 -> 42, 1 -> 43))
+ assertDataStructuresEmpty()
+ }
+
+ test("pipelined shuffle: losing an executor does NOT flip a completed producer to unavailable " +
+ "or resubmit it") {
+ // A completed, consumed pipelined producer whose executor is lost must NOT be
+ // resubmitted (which would hang the streaming writer in awaitTerminationAcks). Because
+ // pipelined
+ // availability is tracked on the stage (monotonic) and not the MapOutputTracker, executor loss
+ // cannot flip isAvailable, so the producer is not resubmitted.
+ val producerRdd = new MyRDD(sc, 2, Nil)
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+ val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ submit(consumerRdd, Array(0, 1))
+ val producerStageId = taskSets.head.stageId
+ val producerStage =
+ scheduler.stageIdToStage(producerStageId).asInstanceOf[ShuffleMapStage]
+ // Producer completes on hostA (both partitions). makeMapStatus("hostA") registers the output
+ // under executor id "hostA-exec" (makeBlockManagerId appends "-exec"), so the ExecutorLost
+ // below -- which loses executor id "hostA-exec" -- actually matches the registered output.
+ complete(taskSets.head, Seq(
+ (Success, makeMapStatus("hostA", 2)),
+ (Success, makeMapStatus("hostA", 2))))
+ assert(producerStage.isAvailable)
+ val taskSetsAfterProducer = taskSets.size
+
+ // Lose the executor that ran the producer ("hostA-exec"). For a regular shuffle this would
+ // strip the outputs and flip isAvailable -> resubmit; for a pipelined shuffle it must be inert.
+ runEvent(ExecutorLost("hostA-exec", ExecutorExited(-100, false, "Container marked as failed")))
+
+ assert(producerStage.isAvailable,
+ "a completed pipelined producer must remain available after executor loss (no tracker strip)")
+ // Pin the underlying completed set directly (not just the derived isAvailable boolean):
+ // executor
+ // loss must not remove any completed partition, so nothing is missing.
+ assert(producerStage.findMissingPartitions() === Seq.empty,
+ "executor loss must not strip a pipelined producer's completed partitions (monotonic)")
+ assert(taskSets.size === taskSetsAfterProducer,
+ "the pipelined producer must NOT be resubmitted on executor loss")
+
+ // The consumer completes normally; no hang, no extra producer attempt.
+ val consumerTaskSet = taskSets.find { ts =>
+ scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd
+ }.get
+ complete(consumerTaskSet, Seq((Success, 42), (Success, 43)))
+ assert(results === Map(0 -> 42, 1 -> 43))
+ assertDataStructuresEmpty()
+ }
+
+ test("pipelined shuffle: binding a live producer stage to a second concurrent job fails fast") {
+ // A pipelined shuffle is a once-through live stream with no retained output, so two concurrent
+ // jobs cannot share one producer stage. While job 0 is still active its producer
+ // stage stays cached in shuffleIdToMapStage bound to job 0; a second concurrent job that reuses
+ // the SAME PipelinedShuffleDependency would bind that live stage to a second jobId, which is
+ // the
+ // forbidden cross-job reuse. Fail fast rather than let job 1 attach to job 0's live stream.
+ // (A sequential re-run is NOT this case: after job 0 finishes, cleanup drops the stage from
+ // shuffleIdToMapStage, so a later job gets a fresh producer bound to only its own job --
+ // exactly how each streaming micro-batch reruns its producer.)
+ val producerRdd = new MyRDD(sc, 2, Nil)
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+
+ // Job 0: submit but leave it active (do not complete the producer or the result stage), so the
+ // producer stage remains cached in shuffleIdToMapStage bound to job 0.
+ val firstConsumer = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ submit(firstConsumer, Array(0, 1))
+ val producerStage =
+ scheduler.stageIdToStage(taskSets.head.stageId).asInstanceOf[ShuffleMapStage]
+ assert(producerStage.jobIds === Set(0))
+ val shuffleStagesBeforeJob1 = scheduler.shuffleIdToMapStage.size
+ val stagesBeforeJob1 = scheduler.stageIdToStage.size
+
+ // Job 1: a second concurrent job reusing the SAME pipelined dependency -> cross-job reuse of a
+ // live producer stage -> fail fast.
+ val secondConsumer = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ val failure = new java.util.concurrent.atomic.AtomicReference[Exception]()
+ val failListener = new JobListener {
+ override def taskSucceeded(index: Int, result: Any): Unit = {}
+ override def jobFailed(exception: Exception): Unit = failure.set(exception)
+ }
+ submit(secondConsumer, Array(0, 1), listener = failListener)
+ assert(failure.get() != null,
+ "binding a live pipelined producer to a second concurrent job must fail the job")
+ assert(failure.get().getMessage.contains("PIPELINED_SHUFFLE_CROSS_JOB_REUSE") ||
+ failure.get().getMessage.contains("reused across jobs"),
+ s"expected a cross-job-reuse error, got: ${failure.get().getMessage}")
+ // Job 0's producer stage was never bound to job 1, and the failed job 1 left no scheduler
+ // residue (the throw happened during stage creation, before any job-1 state was registered).
+ assert(producerStage.jobIds === Set(0))
+ assert(scheduler.shuffleIdToMapStage.size === shuffleStagesBeforeJob1,
+ "the failed cross-job submission must not add a shuffle->stage mapping")
+ assert(scheduler.stageIdToStage.size === stagesBeforeJob1,
+ "the failed cross-job submission must not leave orphaned stages")
+
+ // Job 0 can still complete normally -- the rejected job 1 did not corrupt its shared producer.
+ completeShuffleMapStageSuccessfully(producerStage.id, 0, 2)
+ val job0Consumer = taskSets.find { ts =>
+ scheduler.stageIdToStage(ts.stageId).rdd eq firstConsumer
+ }.get
+ complete(job0Consumer, Seq((Success, 42), (Success, 43)))
+ assert(results === Map(0 -> 42, 1 -> 43))
+ assertDataStructuresEmpty()
+ }
+
+ test("pipelined shuffle: sequential re-run of the same producer is sound (fresh stage per job)") {
+ // Contrast with the concurrent case above: after a job using a pipelined dependency finishes,
+ // cleanup removes its producer from shuffleIdToMapStage, so re-submitting a job on the SAME
+ // dependency creates a fresh producer stage bound to only the new job. This is sound (it is how
+ // a streaming micro-batch reruns) and must NOT trip the cross-job-reuse fail-fast.
+ val producerRdd = new MyRDD(sc, 2, Nil)
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+
+ val consumer0 = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ submit(consumer0, Array(0, 1))
+ complete(taskSets(0), Seq(
+ (Success, makeMapStatus("hostA", 2)),
+ (Success, makeMapStatus("hostB", 2))))
+ complete(taskSets(1), Seq((Success, 42), (Success, 43)))
+ assert(results === Map(0 -> 42, 1 -> 43))
+ assertDataStructuresEmpty()
+ results.clear()
+
+ // Second, sequential job on the same dependency: a fresh producer stage, no fail-fast, no hang.
+ // taskSets is a growing buffer across both jobs, so the second job's task sets are at index 2+.
+ val consumer1 = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ submit(consumer1, Array(0, 1))
+ complete(taskSets(2), Seq(
+ (Success, makeMapStatus("hostA", 2)),
+ (Success, makeMapStatus("hostB", 2))))
+ complete(taskSets(3), Seq((Success, 4), (Success, 5)))
+ assert(results === Map(0 -> 4, 1 -> 5))
+ assertDataStructuresEmpty()
+ }
+
+ test("pipelined shuffle: a FetchFailed on a group member fails the group, not a single-stage " +
+ "resubmit") {
+ // A FetchFailed must fail a pipelined (streaming) query promptly rather than trigger the
+ // base scheduler's single-stage resubmit -> serial recompute -> deadlock. The transient
+ // pipelined shuffle cannot be re-read, and members are co-scheduled, so a lone-stage resubmit
+ // is never valid. Any member's FetchFailed must abort the whole group (-> job abort ->
+ // the caller reruns the batch). Note the base TaskSetManager does NOT count a FetchFailed (it
+ // marks the task successful and zombies the set), so the group-atomic maxTaskFailures=1 lever
+ // does not apply to FetchFailed; the routing must be enforced in the DAGScheduler.
+ val producerRdd = new MyRDD(sc, 2, Nil)
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+ val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ submit(consumerRdd, Array(0, 1))
+ val producerTs = taskSets.head
+ val producerStage =
+ scheduler.stageIdToStage(producerTs.stageId).asInstanceOf[ShuffleMapStage]
+
+ // Producer completes (its outputs are tracked locally on the pipelined stage).
+ complete(producerTs, Seq(
+ (Success, makeMapStatus("hostA", 2)),
+ (Success, makeMapStatus("hostB", 2))))
+ val consumerTs = taskSets.find { ts =>
+ scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd
+ }.get
+ val taskSetsBeforeFetchFailure = taskSets.size
+
+ // A consumer task hits a FetchFailed reading the pipelined shuffle.
+ runEvent(makeCompletionEvent(
+ consumerTs.tasks(0),
+ FetchFailed(makeBlockManagerId("hostA"), pipelinedDep.shuffleId, 0L, 0, 0, "ignored"),
+ null))
+
+ // The job must be failed (group-atomic), and no single stage may be resubmitted: no new task
+ // set is created, and the scheduler is not left waiting to recompute the producer in isolation.
+ // scheduleResubmit posts ResubmitFailedStages on a timer, so drive any pending resubmit and
+ // confirm nothing new is launched.
+ scheduler.resubmitFailedStages()
+ assert(failure != null, "a FetchFailed on a pipelined group member must fail the job")
+ assert(taskSets.size === taskSetsBeforeFetchFailure,
+ "a pipelined group member's FetchFailed must NOT resubmit a single stage")
+ assert(!scheduler.runningStages.exists(_.isInstanceOf[ShuffleMapStage]),
+ "the pipelined producer must not be left running/resubmitted after the group fails")
+ sc.listenerBus.waitUntilEmpty()
+ assertDataStructuresEmpty()
+ }
+
+ test("pipelined shuffle: a consumer result task throwing aborts the whole group") {
+ // Defense-in-depth for the group-atomic model at the DAGScheduler layer: a CONSUMER (result)
+ // task failing must tear down the whole group, not just its own stage. Result and map tasks
+ // share handleFailedTask / effectiveMaxTaskFailures=1, so the first consumer-task exception
+ // makes the TaskSetManager abort the set (delivered here as TaskSetFailed, exactly as a
+ // maxTaskFailures=1 abort surfaces to the DAGScheduler), which must abort the group and tear
+ // down the still-running producer -- the caller then reruns the batch. Complements the
+ // TaskSetManager-layer maxTaskFailures=1 tests and the producer-failure drop test.
+ val producerRdd = new MyRDD(sc, 2, Nil)
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+ val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ submit(consumerRdd, Array(0, 1))
+ val consumerTs = taskSets.find { ts =>
+ scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd
+ }.get
+ val producerStage = scheduler.stageIdToStage(taskSets.head.stageId)
+ // The producer is STILL RUNNING (co-scheduled, not yet completed) when the consumer fails --
+ // this is what gives the assertion teeth: group-atomic teardown must reach a running producer.
+ assert(scheduler.runningStages.contains(producerStage),
+ "the pipelined producer must be co-scheduled and running alongside the consumer")
+
+ // The consumer's task set aborts on the first task exception (maxTaskFailures=1 for a pipelined
+ // member). This must fail the whole group and tear down the still-running producer.
+ failed(consumerTs, "consumer result task threw")
+ assert(failure != null, "a consumer task failure must fail the group's job")
+ assert(!scheduler.runningStages.contains(producerStage),
+ "the still-running pipelined producer must be torn down when the consumer fails the group")
+ sc.listenerBus.waitUntilEmpty()
+ assertDataStructuresEmpty()
+ }
+
+ // ==========================================================================================
+ // Group-atomic rerun resets per-partition commit authorization
+ // ==========================================================================================
+
+ test("pipelined shuffle: a group rerun resets per-partition commit authorization") {
+ // A pipelined group is atomic, so a failure reruns the WHOLE group -- including a result
+ // stage whose tasks already succeeded and committed. Those committed partitions are rerun and
+ // must be allowed to commit again. OutputCommitCoordinator permanently denies re-commit for a
+ // committed partition (a Success clears nothing; keyed by stage id). The rerun stays correct
+ // two ways, both asserted here: (b) the group teardown runs the committed result stage through
+ // markStageAsFinished -> stageEnd, clearing its committer state (so no stale authorization
+ // survives); and (a) the caller's rerun is a NEW job whose stages get FRESH stage ids, so the
+ // coordinator has no prior committer for them regardless.
+ val producerRdd = new MyRDD(sc, 2, Nil)
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+ val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ val failure = new java.util.concurrent.atomic.AtomicReference[Exception]()
+ val failListener = new JobListener {
+ override def taskSucceeded(index: Int, result: Any): Unit = {}
+ override def jobFailed(exception: Exception): Unit = failure.set(exception)
+ }
+ submit(consumerRdd, Array(0, 1), listener = failListener)
+ // Locate producer/consumer defensively by RDD identity (not task-set order).
+ val producerTaskSet =
+ taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq producerRdd).get
+ val consumerTaskSet =
+ taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd).get
+ val firstConsumerStageId = consumerTaskSet.stageId
+
+ // Both stages register OutputCommitCoordinator state at submission (stageStart), so the
+ // coordinator holds per-stage commit-authorization state that a rerun must not inherit. (This
+ // asserts stage state exists to be reset; it does not claim a committer was authorized -- that
+ // needs a live canCommit RPC, which the mock backend does not drive.)
+ assert(!scheduler.outputCommitCoordinator.isEmpty,
+ "the coordinator must hold commit-authorization state for the submitted group's stages")
+
+ // The producer now fails -> group-atomic failure -> the whole group is torn down and the job
+ // fails; the caller will rerun the batch as a new job.
+ failed(producerTaskSet, "producer blew up")
+ assert(failure.get() != null, "the group must fail atomically when the producer fails")
+ // Teardown ran the group's stages through markStageAsFinished -> stageEnd, clearing their
+ // coordinator state: the commit-authorization state does not survive the failed attempt.
+ // (isEmpty here proves stageEnd reached every stage the teardown covered.)
+ assert(scheduler.outputCommitCoordinator.isEmpty,
+ "group teardown must reset per-partition commit authorization; none may survive")
+ assertDataStructuresEmpty()
+
+ // The caller reruns the batch as a NEW job on the same dependency. Its stages get fresh ids, so
+ // the coordinator has no prior committer, and the rerun's partition 0 can commit again. Only
+ // the
+ // task sets submitted from here on belong to the rerun (earlier ones' stages were cleaned up,
+ // so
+ // look them up defensively).
+ val taskSetsBeforeRerun = taskSets.size
+ val rerunConsumer = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ submit(rerunConsumer, Array(0, 1))
+ val rerunTaskSets = taskSets.drop(taskSetsBeforeRerun)
+ val rerunConsumerTaskSet = rerunTaskSets.find { ts =>
+ scheduler.stageIdToStage.get(ts.stageId).exists(_.rdd eq rerunConsumer)
+ }.get
+ val rerunProducerTaskSet = rerunTaskSets.find { ts =>
+ scheduler.stageIdToStage.get(ts.stageId).exists(_.rdd eq producerRdd)
+ }.get
+ assert(rerunConsumerTaskSet.stageId != firstConsumerStageId,
+ "the rerun's result stage must get a fresh stage id (fresh coordinator state)")
+ complete(rerunProducerTaskSet, Seq(
+ (Success, makeMapStatus("hostA", 2)),
+ (Success, makeMapStatus("hostB", 2))))
+ complete(rerunConsumerTaskSet, Seq((Success, 7), (Success, 8)))
+ assert(results === Map(0 -> 7, 1 -> 8), "the rerun must complete, re-committing its partitions")
+ assertDataStructuresEmpty()
+ }
+
+ test("pipelined shuffle: a buffered consumer success is dropped on group abort, and its " +
+ "TaskEnd is flushed so the stage is not leaked as running") {
+ // A consumer's successful task has its whole completion event buffered while its producer runs
+ // (coarse model) -- no TaskEnd yet. If the group is then torn down before the producer finishes
+ // (here: the consumer's OTHER task hits a FetchFailed, which aborts the group), the buffered
+ // success's result must NOT be applied, but its TaskEnd IS flushed on the drop path. A listener
+ // that tracks active tasks (e.g. AppStatusListener, which removes a stage once activeTasks hits
+ // 0) therefore does not leak the consumer stage as perpetually running: the success's TaskStart
+ // was delivered, and its matching TaskEnd arrives at teardown even though the result is
+ // dropped.
+ val endedTaskIds = new java.util.concurrent.ConcurrentHashMap[Long, Boolean]()
+ val recordingListener = new SparkListener {
+ override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit =
+ endedTaskIds.put(taskEnd.taskInfo.taskId, true)
+ }
+ sc.addSparkListener(recordingListener)
+ try {
+ val producerRdd = new MyRDD(sc, 2, Nil)
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+ val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ val failure = new java.util.concurrent.atomic.AtomicReference[Exception]()
+ val failListener = new JobListener {
+ override def taskSucceeded(index: Int, result: Any): Unit = results.put(index, result)
+ override def jobFailed(exception: Exception): Unit = failure.set(exception)
+ }
+ submit(consumerRdd, Array(0, 1), listener = failListener)
+ // Identify the co-scheduled consumer (ResultStage over consumerRdd) and its still-running
+ // pipelined producer by RDD identity rather than task-set order.
+ val consumerTaskSet = taskSets.find { ts =>
+ scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd
+ }.get
+ assert(scheduler.dependentStageMap.size === 1,
+ "the consumer must be co-scheduled with a running producer (a deferral must exist)")
+
+ // Consumer partition 0 succeeds. Its whole completion event is buffered (no job result and no
+ // TaskEnd yet). Give it a known taskId so we can assert the drop path flushes it later.
+ val bufferedTaskId = 7007L
+ runEvent(makeCompletionEvent(consumerTaskSet.tasks(0), Success, 42,
+ taskInfo = createFakeTaskInfoWithId(bufferedTaskId)))
+ sc.listenerBus.waitUntilEmpty()
+ assert(!endedTaskIds.containsKey(bufferedTaskId),
+ "the buffered consumer success must not emit its TaskEnd while the producer still runs")
+ assert(results.isEmpty, "the consumer's job result must still be deferred (producer running)")
+ assert(scheduler.dependentStageMap.get(scheduler.stageIdToStage(consumerTaskSet.stageId))
+ .exists(_.delayedTaskCompletionEvents.nonEmpty),
+ "the consumer's completion event must be buffered while its producer runs")
+
+ // The consumer's OTHER task now hits a FetchFailed. For a pipelined group member this aborts
+ // the whole group (group-atomic failure) WITHOUT the producer ever finishing. The buffered
+ // success is dropped (its result must not be applied), but its TaskEnd is flushed on the drop
+ // path so active-task-tracking listeners see the task finish.
+ runEvent(makeCompletionEvent(
+ consumerTaskSet.tasks(1),
+ FetchFailed(makeBlockManagerId("hostA"), pipelinedDep.shuffleId, 0L, 0, 0, "ignored"),
+ null))
+ scheduler.resubmitFailedStages()
+ assert(failure.get() != null, "the job must fail")
+
+ sc.listenerBus.waitUntilEmpty()
+ assert(results.isEmpty,
+ "the deferred consumer success must be dropped on abort, not applied as a result")
+ assert(endedTaskIds.containsKey(bufferedTaskId),
+ "the buffered consumer success's TaskEnd must be flushed on the drop path, so the stage " +
+ "is not leaked as perpetually running")
+ assertDataStructuresEmpty()
+ } finally {
+ sc.removeSparkListener(recordingListener)
+ }
+ }
+
+ test("pipelined shuffle: a barrier producer stage is rejected") {
+ // A barrier stage exposes output only after a global sync, incompatible with incremental reads.
+ val producerRdd = new MyRDD(sc, 2, Nil).barrier().mapPartitions(iter => iter)
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+ val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ assertPipelinedUnsupported(submitAndCaptureFailure(consumerRdd, Array(0, 1)), "barrier")
+ assertDataStructuresEmpty()
+ }
+
+ test("pipelined shuffle: a statically-indeterminate producer is rejected") {
+ // Indeterminate output's recovery is stage rollback-and-recompute, which a group never performs
+ // (so it never applies); reject rather than carry dead machinery.
+ val producerRdd = new MyRDD(sc, 2, Nil, indeterminate = true)
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+ val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ assertPipelinedUnsupported(submitAndCaptureFailure(consumerRdd, Array(0, 1)), "indeterminate")
+ assertDataStructuresEmpty()
+ }
+
+ test("pipelined shuffle: a reliable RDD checkpoint in a PRODUCER's chain is rejected") {
+ // A reliable checkpoint writes a durable, lineage-truncated snapshot -> reintroduces cross-time
+ // reuse of a transient edge and needs a post-success recompute of the vanished input. Rejected
+ // by walking the producer's within-stage chain for a ReliableRDDCheckpointData. Keyed on
+ // checkpointData (not isCheckpointed): the write has not happened yet at group-creation time.
+ withTempDir { dir =>
+ sc.setCheckpointDir(dir.getCanonicalPath)
+ val checkpointableRdd = new MyCheckpointRDD(sc, 2, Nil)
+ checkpointableRdd.checkpoint() // sets checkpointData to a ReliableRDDCheckpointData
+ assert(checkpointableRdd.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]]))
+ val pipelinedDep = new PipelinedShuffleDependency(checkpointableRdd, new HashPartitioner(2))
+ val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ assertPipelinedUnsupported(
+ submitAndCaptureFailure(consumerRdd, Array(0, 1)), "reliable RDD checkpoint")
+ assertDataStructuresEmpty()
+ }
+ }
+
+ test("pipelined shuffle: a reliable RDD checkpoint in a CONSUMER's chain is rejected") {
+ // The rejection must cover a consumer member's chain too, not just the producer's: a consumer's
+ // transient input IS the pipelined shuffle, so a reliable checkpoint there would re-read the
+ // vanished stream on recompute. Purely all-pipelined shape (no regular shuffle -- those are
+ // rejected separately): producer --pipelined--> consumer(checkpointed, result). The consumer
+ // reads the pipelined shuffle and its within-stage chain carries the reliable checkpoint.
+ withTempDir { dir =>
+ sc.setCheckpointDir(dir.getCanonicalPath)
+ val producerRdd = new MyRDD(sc, 2, Nil)
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+ // The consumer RDD reads the pipelined shuffle AND is reliably checkpointed.
+ val consumerRdd = new MyCheckpointRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ consumerRdd.checkpoint()
+ assert(consumerRdd.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]]))
+ assertPipelinedUnsupported(
+ submitAndCaptureFailure(consumerRdd, Array(0, 1)), "reliable RDD checkpoint")
+ assertDataStructuresEmpty()
+ }
+ }
+
+ test("pipelined shuffle: a reliable checkpoint DOWNSTREAM in the consumer stage (not on the " +
+ "reading RDD) is still rejected") {
+ // Coverage for a checkpoint that is NOT on the RDD reading the pipelined shuffle, but on a
+ // narrow-dep child of it WITHIN the same consumer stage. Purely all-pipelined shape:
+ // producer --pipelined--> reads --(narrow)--> checkpointed (result).
+ // `reads` and `checkpointed` are one stage. Rooting the check at the pipelined-reading RDD and
+ // walking parents would MISS this (the checkpoint is downstream of the read); the check must
+ // instead recognize that `checkpointed`'s own within-stage chain reads a pipelined shuffle.
+ withTempDir { dir =>
+ sc.setCheckpointDir(dir.getCanonicalPath)
+ val producerRdd = new MyRDD(sc, 2, Nil)
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+ val readsRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ // A narrow-dep child of the reading RDD, in the SAME stage, that is reliably checkpointed,
+ // and is the result RDD (no regular shuffle after -- that would be separately rejected).
+ val checkpointedRdd = new MyCheckpointRDD(sc, 2, List(new OneToOneDependency(readsRdd)))
+ checkpointedRdd.checkpoint()
+ assert(checkpointedRdd.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]]))
+ assertPipelinedUnsupported(
+ submitAndCaptureFailure(checkpointedRdd, Array(0, 1)), "reliable RDD checkpoint")
+ assertDataStructuresEmpty()
+ }
+ }
+
+ test("pipelined shuffle: a producer feeding more than one consumer (fan-out) is rejected") {
+ // 1:N fan-out is deferred to a later version and rejected up front here. Fan-out is
+ // detected at the RDD level -- two DISTINCT RDDs listing the same pipelined shuffle as a
+ // dependency -- so it is expressible in an all-pipelined job without any regular shuffle: two
+ // consumer RDDs both read the same pipelined producer, unioned by a narrow dependency into the
+ // result. checkPipelinedGroupsSupportedInRDDGraph counts 2 distinct consumers for the shuffle.
+ val producerRdd = new MyRDD(sc, 2, Nil)
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+ val consumerA = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ val consumerB = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ val union = new MyRDD(sc, 2,
+ List(new OneToOneDependency(consumerA), new OneToOneDependency(consumerB)),
+ tracker = mapOutputTracker)
+ assertPipelinedUnsupported(
+ submitAndCaptureFailure(union, Array(0, 1)), "more than one consumer")
+ assertDataStructuresEmpty()
+ }
+
+ // Resource-profile rejection. The gang slot check measures capacity against the DEFAULT
+ // resource profile, so the whole group must run on the default profile; any member with an
+ // explicit non-default profile is rejected. The three shapes below must all be rejected; the
+ // uniform-non-default and non-default-plus-default cases in particular would each pass a
+ // "more than one DISTINCT profile" check yet still be admitted against the wrong (default) pool.
+ private def rpWithCores(cores: Int, cpus: Int): ResourceProfile =
+ new ResourceProfileBuilder()
+ .require(new ExecutorResourceRequests().cores(cores))
+ .require(new TaskResourceRequests().cpus(cpus)).build()
+
+ test("pipelined shuffle: a group with two distinct non-default resource profiles is rejected") {
+ val producerRdd = new MyRDD(sc, 2, Nil).withResources(rpWithCores(4, 1))
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+ val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ .withResources(rpWithCores(8, 2))
+ assertPipelinedUnsupported(
+ submitAndCaptureFailure(consumerRdd, Array(0, 1)), "non-default resource profile")
+ assertDataStructuresEmpty()
+ }
+
+ test("pipelined shuffle: a group uniformly on one non-default resource profile is rejected") {
+ // Both members share ONE non-default profile -- a "distinct profile count" of 1, which a
+ // more-than-one-distinct-profile check would wrongly admit, then measure against the default
+ // profile's capacity (the wrong pool). Must be rejected: the whole group is off the default RP.
+ val rp = rpWithCores(4, 1)
+ val producerRdd = new MyRDD(sc, 2, Nil).withResources(rp)
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+ val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ .withResources(rp)
+ assertPipelinedUnsupported(
+ submitAndCaptureFailure(consumerRdd, Array(0, 1)), "non-default resource profile")
+ assertDataStructuresEmpty()
+ }
+
+ test("pipelined shuffle: a group mixing a non-default profile with the default is rejected") {
+ // The producer carries an explicit non-default profile; the consumer is left on the default.
+ // The set of EXPLICIT profiles is again size 1, so a distinct-explicit-profile check would miss
+ // it, but the group genuinely spans the non-default and default profiles. Must be rejected.
+ val producerRdd = new MyRDD(sc, 2, Nil).withResources(rpWithCores(4, 1))
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+ val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ assertPipelinedUnsupported(
+ submitAndCaptureFailure(consumerRdd, Array(0, 1)), "non-default resource profile")
+ assertDataStructuresEmpty()
+ }
+
+ test("regular shuffle idioms are NOT rejected (inertness of the pipelined fail-fast)") {
+ // The pipelined fail-fast checks (indeterminate producer, reliable-checkpoint-in-chain) must be
+ // inert for a job with NO pipelined dependency: a regular shuffle whose producer is BOTH
+ // indeterminate AND reliably checkpointed -- the two idioms most likely to false-fire -- must
+ // run exactly as before. (Actually exercise both, not just the flag.)
+ withTempDir { dir =>
+ sc.setCheckpointDir(dir.getCanonicalPath)
+ val producerRdd = new MyCheckpointRDD(sc, 2, Nil, indeterminate = true)
+ producerRdd.checkpoint()
+ assert(producerRdd.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]]))
+ assert(producerRdd.outputDeterministicLevel == DeterministicLevel.INDETERMINATE)
+ val regularDep = new ShuffleDependency(producerRdd, new HashPartitioner(2))
+ val consumerRdd = new MyRDD(sc, 2, List(regularDep), tracker = mapOutputTracker)
+ submit(consumerRdd, Array(0, 1))
+ // The job proceeds normally (producer stage submitted), i.e. NOT failed by a pipelined check.
+ assert(failure === null, "a regular shuffle must not be rejected by the pipelined fail-fast")
+ completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2)
+ complete(taskSets(1), Seq((Success, 42), (Success, 43)))
+ assert(results === Map(0 -> 42, 1 -> 43))
+ assertDataStructuresEmpty()
+ }
+ }
+
+ test("regular job on a non-default resource profile is NOT rejected (RP check is pipelined)") {
+ // The resource-profile rejection is not keyed on a pipelined dependency, so it must run ONLY
+ // for a job that has one (handleJobSubmitted gates checkPipelinedGroupsSupportedInRDDGraph on
+ // hasPipelined). A perfectly ordinary job that merely attaches a non-default profile via
+ // RDD.withResources -- a GA stage-level-scheduling feature -- has NO pipelined dependency and
+ // must run untouched. Without the gate the whole graph walk fires and rejects it with
+ // PIPELINED_SHUFFLE_UNSUPPORTED, a regression on plain withResources jobs.
+ val rdd = new MyRDD(sc, 2, Nil).withResources(rpWithCores(4, 2))
+ // A non-default profile drives submitMissingTasks through addPySparkConfigsToProperties, which
+ // needs a non-null Properties; pass one (the default submit() overload leaves it null).
+ submit(rdd, Array(0, 1), properties = new Properties())
+ assert(failure === null,
+ "a regular job with a non-default resource profile must not be rejected by the pipelined " +
+ "fail-fast")
+ assert(taskSets.nonEmpty, "the job must proceed to task submission, not be failed up front")
+ complete(taskSets(0), Seq((Success, 42), (Success, 43)))
+ assert(results === Map(0 -> 42, 1 -> 43))
+ assertDataStructuresEmpty()
+ }
+
+ private def submitAndCaptureFailure(finalRdd: RDD[_], partitions: Array[Int]): Exception = {
+ val failure = new java.util.concurrent.atomic.AtomicReference[Exception]()
+ val failListener = new JobListener {
+ override def taskSucceeded(index: Int, result: Any): Unit = {}
+ override def jobFailed(exception: Exception): Unit = failure.set(exception)
+ }
+ submit(finalRdd, partitions, listener = failListener)
+ failure.get()
+ }
+
+ private def assertPipelinedUnsupported(failure: Exception, reasonSubstring: String): Unit = {
+ assert(failure != null, "the job must fail fast on the unsupported pipelined idiom")
+ val msg = failure.getMessage
+ assert(msg.contains("PIPELINED_SHUFFLE_UNSUPPORTED") || msg.contains("unsupported feature"),
+ s"expected a PIPELINED_SHUFFLE_UNSUPPORTED error, got: $msg")
+ assert(msg.contains(reasonSubstring),
+ s"expected reason to mention '$reasonSubstring', got: $msg")
+ }
+
}
class DAGSchedulerAbortStageOffSuite extends DAGSchedulerSuite {
diff --git a/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala
index 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) {
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org