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