[WIP] [SPARK-XXXXX][CORE] Concurrently schedule pipelined-shuffle stage groups in the DAGScheduler#57341
Open
jerrypeng wants to merge 35 commits into
Open
[WIP] [SPARK-XXXXX][CORE] Concurrently schedule pipelined-shuffle stage groups in the DAGScheduler#57341jerrypeng wants to merge 35 commits into
jerrypeng wants to merge 35 commits into
Conversation
Introduce PipelinedShuffleDependency, a ShuffleDependency subtype that declares a shuffle's output can be read incrementally -- a consumer stage may begin reading while the producer is still running. This is the first-class marker a later DAGScheduler change will use to co-schedule the producer and consumer (a 'pipelined group') and to select an incremental shuffle implementation. This PR adds only the type. On its own it behaves exactly like its parent ShuffleDependency -- code that matches ShuffleDependency continues to treat it as an ordinary materialized shuffle -- so it changes no existing behavior. The concurrent-scheduling and incremental-shuffle behavior are added separately. The parent's checksumMismatchFullRetryEnabled / checksumMismatchQueryLevelRoll- backEnabled params are intentionally not exposed, so they stay false: their stage-level recompute-and-rerun is moot for a pipelined group (any failure aborts the whole group and the caller reruns from scratch), and would also conflict with a consumer that has already read the output incrementally. Tests: PipelinedShuffleDependency is-a ShuffleDependency, is distinguishable from an ordinary dependency, registers its shuffle with a distinct shuffleId, keeps the checksum retry/rollback flags false, and forwards non-default constructor args (aggregator, mapSideCombine) to the parent correctly. Co-authored-by: Isaac
…eManager by dependency type Introduce spark.shuffle.manager.incremental: when set, SparkEnv instantiates a second ShuffleManager (a peer of the default spark.shuffle.manager, not a wrapper) and SparkEnv.shuffleManagerFor(dep) routes each shuffle to one of the two by dependency type -- a PipelinedShuffleDependency to the incremental manager, every other ShuffleDependency to the default. Routing is a pure function of the dependency, so the driver (registerShuffle) and executors (getReader/getWriter) agree without shared state; each manager only ever sees its own handle type, so no wrapping handle is needed. SparkEnv.get.shuffleManager remains the plain default manager, so nothing is installed in front of it.
…uching the default-manager check Restructure initializeStreamingShuffleOutputTracker so the pre-existing default-manager branch (which matches StreamingShuffleManager or MultiShuffleManager) is left untouched, and the new incremental-manager case (spark.shuffle.manager.incremental == StreamingShuffleManager) is handled by a separate early return. Extract the tracker construction into createStreamingShuffleOutputTracker so both paths share it. This keeps MultiShuffleManager out of this PR's diff entirely: the reference that remains is master's default-slot check, which this routing PR does not modify. Behavior is unchanged (routing suite tracker-init tests still pass). Co-authored-by: Isaac
…callers need no ShuffleManager Block-by-id resolution always uses the default manager (a pipelined shuffle is served out-of-band and produces no block-manager blocks). Add a SparkEnv.shuffleBlockResolver accessor that returns defaultShuffleManager.shuffleBlockResolver, and route the block-serving paths through it instead of SparkEnv.get.defaultShuffleManager.shuffleBlockResolver: - BlockManager: resolver accesses (getBlockData, merged blocks, migratableResolver, corruption diagnosis) go through a private shuffleBlockResolver accessor that preserves the deferred-init wait and the injected-manager (test) path. The manager reference is kept only for the unsupported-resolver error. - BlockManagerMasterEndpoint: the ESS-cleanup path uses the resolver accessor directly. Centralizes the "block resolution uses the default manager" invariant in one place. Updates the BlockManagerSuite mock to stub the new accessor alongside defaultShuffleManager. Co-authored-by: Isaac
…elinedShuffle by kind A ShuffleManager's type now declares what kind of shuffle it implements, instead of the kind being implicit in whether shuffleBlockResolver throws: - BlockingShuffle: materializes output as block-manager-addressed blocks served through a ShuffleBlockResolver (reads, push-merge, decommission migration). shuffleBlockResolver moves off ShuffleManager onto this trait. SortShuffleManager extends it. - PipelinedShuffle: output is read incrementally and served out-of-band, so there is no ShuffleBlockResolver. StreamingShuffleManager extends it and drops its throwing resolver stub. Callers that resolve blocks now match on BlockingShuffle rather than assuming every manager has a resolver: - SparkEnv.shuffleBlockResolver returns Option[ShuffleBlockResolver] (None when the default manager is not a BlockingShuffle). - BlockManager / BlockManagerMasterEndpoint handle the None case; block-serving paths require a BlockingShuffle default (a clear error otherwise, matching the previous throw). - ShuffleWriteProcessor's push-merge path matches BlockingShuffle before touching the resolver. - SparkCoreErrors.unexpectedShuffleBlockWithUnsupportedResolverError takes the resolver directly. MultiShuffleManager stays a plain ShuffleManager (it is slated for deprecation in favor of per-dependency routing) and keeps its own resolver method for the SortShuffleManager it delegates to. New tests assert the manager-kind types and that SparkEnv.shuffleBlockResolver is defined only for a blocking default manager. Co-authored-by: Isaac
…'s shuffle block resolver Deprecate MultiShuffleManager: the cluster-level single-slot approach is superseded by per-dependency routing (spark.shuffle.manager + spark.shuffle.manager.incremental). As a plain ShuffleManager (not a BlockingShuffle), it no longer participates in block-by-id resolution, ESS cleanup, or push-based merge when configured as the default manager -- this is a documented, accepted trade-off for a manager on its way out. Also restore BlockManager.shuffleBlockResolver to a cached lazy val (it had become a def in the trait-split change, recomputing the init-wait and type match on every block access in hot paths like getBlockData). Co-authored-by: Isaac
…ency routing MultiShuffleManager was the cluster-level single-slot way to run streaming and regular shuffles together. Per-dependency routing (spark.shuffle.manager for regular/blocking shuffles and spark.shuffle.manager.incremental for pipelined ones, each routed by dependency type) supersedes it, so remove the class and its suite. This also removes the only in-tree ShuffleManager that was neither a BlockingShuffle nor a PipelinedShuffle: the streaming-tracker init no longer needs to special-case it, and the default manager is now always a BlockingShuffle (built-in) unless a user configures a custom one. Co-authored-by: Isaac
…al ShuffleManager, simplify ESS resolver lookup - Rename BlockingShuffle -> BlockingShuffleManager and PipelinedShuffle -> PipelinedShuffleManager for consistency with the ShuffleManager they extend. - Seal ShuffleManager so it cannot be implemented directly; a manager declares its kind by extending BlockingShuffleManager or PipelinedShuffleManager. - BlockManagerMasterEndpoint's ESS-cleanup path calls SparkEnv.get.shuffleBlockResolver directly instead of a local wrapper val. Co-authored-by: Isaac
…ckingShuffleManager) Per review, keep MultiShuffleManager.scala and its suite rather than deleting them; the new per-dependency routing simply does not reference the class. Since ShuffleManager is now sealed, MultiShuffleManager extends BlockingShuffleManager (it resolves blocks via the inner SortShuffleManager it delegates regular shuffles to). Restore the SparkEnv streaming-tracker-init recognition of a MultiShuffleManager default, and drop a redundant comment in StreamingShuffleManager. Co-authored-by: Isaac
…drop dead endpoint param Address three review comments on the dependency-typed routing PR: - Do not seal ShuffleManager (reverses the earlier seal). sealed would restrict the third-party ShuffleManager extension point from SPARK-45762 (user-jar managers loaded reflectively), and being frontend-only it does not even enforce that against Java or reflective implementations -- its only effect is forcing in-tree subtypes into one file. The BlockingShuffleManager / PipelinedShuffleManager split stays; the "declare your kind" contract is instead enforced by per-slot startup validation (a follow-up commit). - Remove the now-dead _shuffleManager constructor parameter (and its import) from BlockManagerMasterEndpoint: ESS cleanup resolves blocks via SparkEnv.get.shuffleBlockResolver, so nothing consumes it. Drop the argument at all call sites (SparkEnv and three test suites). - Fix the SPARK-45762 test in SparkSubmitSuite. Its user-jar ShuffleManager is compiled by javac at test runtime, so it is not covered by Test/compile; moving shuffleBlockResolver onto BlockingShuffleManager left its @OverRide dangling. The generated class now implements BlockingShuffleManager (its body already stubs every required method), which also exercises a reflectively-loaded BlockingShuffleManager from a foreign classpath. Co-authored-by: Isaac
…by kind, not a default Per review, do not model one manager as "the default" -- that hides the blocking-vs-pipelined distinction the routing is built on, and it left a latent sharp edge: SparkEnv.shuffleBlockResolver resolved through the default slot only, so a blocking manager configured in the incremental slot would have had unreachable blocks. SparkEnv now holds two managers keyed by kind: - blockingShuffleManager (spark.shuffle.manager): a BlockingShuffleManager that serves all regular, materialized shuffles and owns block-by-id resolution. shuffleBlockResolver always resolves through it -- the single, well-typed source for reads, push-merge, and decommission migration. - pipelinedShuffleManager (spark.shuffle.manager.incremental): an optional PipelinedShuffleManager that serves pipelined dependencies out-of-band. initializeShuffleManager validates each slot's kind and fails fast otherwise: a non-blocking manager (including a pipelined-only or kind-less one) is rejected from the default slot, and a blocking manager from the incremental slot. This closes the unreachable-blocks case by construction and is the runtime enforcement of "declare your kind" that replaces the reverted `sealed`. A MultiShuffleManager is still accepted as the blocking manager -- it is blocking-kind and resolves blocks via its inner SortShuffleManager -- so it keeps working as the single-slot streaming+sort manager while routing continues to key on dependency type. The streaming output tracker is initialized when the incremental manager is a StreamingShuffleManager, or the blocking manager is a MultiShuffleManager; the old "StreamingShuffleManager as the default manager" case is gone, since a bare StreamingShuffleManager is pipelined and rejected from the default slot. ShuffleExchangeExec's sort-shuffle detection reads blockingShuffleManager. Routing tests split the recording test double by kind and cover both fail-fast paths. Co-authored-by: Isaac
…after kind validation The kind-based slot validation added to SparkEnv rejects a non-blocking manager in the spark.shuffle.manager (default) slot. The streaming test suites predate the routing model and configured StreamingShuffleManager -- a PipelinedShuffleManager -- as the default, so they now throw IllegalArgumentException at SparkContext startup. These suites are out-of-diff from the validation change, so they broke CI without failing the routing tests. Move StreamingShuffleManager to the incremental slot (spark.shuffle.manager.incremental), which is the correct slot for a pipelined manager and still initializes the StreamingShuffleOutputTracker the reader/writer constructors assert. StreamingShuffleSuite's assertions that the manager is streaming now check pipelinedShuffleManager instead of the (blocking) default. The tracker-init suite gains a slot-aware helper and a case that a MultiShuffleManager default also initializes the tracker. Co-authored-by: Isaac
…ing; add binding policy
Make the built-in streaming shuffle manager the default for spark.shuffle.manager.incremental, the
same way "sort" is the default for spark.shuffle.manager:
- Register "streaming" as a short alias for StreamingShuffleManager in ShuffleManager.resolveShortName
(alongside "sort"/"tungsten-sort"), so the incremental slot accepts the same style of short names.
- Change spark.shuffle.manager.incremental from optional to createWithDefault("streaming"). An empty
value explicitly disables the incremental slot, in which case a pipelined dependency falls back to
the blocking manager. SparkEnv resolves the (now always-present) config, treating empty as
disabled, and the streaming output tracker initializes based on the instantiated pipelined manager
rather than re-reading the config.
Also add .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) to the config, required by
SparkConfigBindingPolicySuite for any new config (a shuffle manager is a physical-execution choice
that does not change how a view/UDF body resolves). Matches the sibling streaming configs.
Tests: the routing suite's fallback and pre-init cases now set the incremental slot empty to
exercise the "no pipelined manager" path explicitly; the tracker-init suite sets the incremental
slot for every case (empty disables it) and drops the case that put a blocking manager in the
incremental slot (now rejected by the kind validation, and already covered by the routing suite's
non-streaming pipelined double).
Co-authored-by: Isaac
…mirroring the blocking one Per review, model the pipelined manager exactly like the blocking manager: it has a built-in default implementation (the streaming shuffle manager), so it is always created rather than optional. - _pipelinedShuffleManager and the pipelinedShuffleManager accessor are now a plain PipelinedShuffleManager (not Option), initialized in initializeShuffleManager the same way as the blocking manager -- resolving spark.shuffle.manager.incremental (default "streaming") and requiring a PipelinedShuffleManager. There is no "disabled" incremental slot. - shuffleManagerFor no longer falls back to the blocking manager for a pipelined dependency: a pipelined dependency is always served by the pipelined manager. - unregister/stop null-guard both managers (both are null until the deferred init runs), matching the blocking manager's handling. Also address the tracker-init review comments: consolidate the streaming-incremental and MultiShuffleManager-default conditions into a single check, and add a TODO to remove the MultiShuffleManager clause once MultiShuffleManager is gone. Tests: drop the now-impossible "incremental disabled -> fallback" and "tracker disabled" cases; add a case that the tracker initializes by default (streaming incremental manager); the pre-init no-op test nulls both manager fields. Co-authored-by: Isaac
…lity, consistent tracker check Address three non-blocking review comments on the dependency-typed routing PR: - Tighten the four new SparkEnv routing accessors (blockingShuffleManager, pipelinedShuffleManager, shuffleBlockResolver, shuffleManagerFor) to private[spark]. They are internal routing plumbing returning private[spark] manager types, and every caller is in-tree under org.apache.spark, so this matches the stated intent and removes a public-method-returns-package-private-type leak on the @DeveloperAPI SparkEnv class. The pre-existing shuffleManager accessor stays public (now @deprecated) since tightening a shipped @DeveloperAPI accessor would itself be a breaking change. - In initializeStreamingShuffleOutputTracker, inspect the instantiated blocking manager (_blockingShuffleManager.isInstanceOf[MultiShuffleManager]) instead of re-reading the config class name, so it is symmetric with the sibling incrementalIsStreaming check and keys on the object actually created. - Fix an inaccurate comment in BlockManagerMasterEndpoint: shuffleBlockResolver is None only before the manager is initialized, not "when the default manager is not blocking" -- the kind validation in initializeShuffleManager makes a non-blocking default manager impossible. Co-authored-by: Isaac
…rom @DeveloperAPI SparkEnv SparkEnv is @DeveloperAPI, so its Scala doc is rendered into Java API docs by unidoc. The new routing accessors documented their return types with [[BlockingShuffleManager]] / [[PipelinedShuffleManager]] / [[ShuffleBlockResolver]] / [[ShuffleManager]], which are all private[spark] and therefore have no Java doc page. Scaladoc [[X]] becomes javadoc {@link X}, and javadoc fails with "reference not found" for a link to a non-public type, breaking the Javaunidoc / doc build. Render those references as backtick code spans instead of doc links (the same idiom the deprecated shuffleManager accessor already uses). The public @DeveloperAPI types PipelinedShuffleDependency and ShuffleDependency keep their [[...]] links, since those resolve. Co-authored-by: Isaac
…nregisterShuffle contract The id-only RemoveShuffle cleanup path (SparkEnv.unregisterShuffleFromAllManagers) cannot tell which manager owns a shuffle, so it now broadcasts unregisterShuffle to every configured manager -- including ones that never handled the shuffle. That relies on an unknown/already-removed id being a harmless no-op, which the built-in managers honor but the ShuffleManager interface never promised. Per review, make it an explicit part of the contract: unregisterShuffle must treat an unknown or already-removed shuffleId as an idempotent no-op (return false) rather than throwing or mutating unrelated state. Documentation-only change to the trait scaladoc. Co-authored-by: Isaac
jerrypeng
force-pushed
the
stack/pipelined-shuffle-pr3-scheduling
branch
from
July 17, 2026 21:54
9aefc6d to
f7a6c5e
Compare
jerrypeng
changed the title
…cy's active behavior; fix ESS-coexistence note Two doc follow-ups from cloud-fan's review: - PipelinedShuffleDependency scaladoc said it "only declares the capability" and "behaves exactly like its parent". That is stale: construction now routes shuffle registration to the incremental manager (via SparkEnv.shuffleManagerFor) and unconditionally disables push-based shuffle merge (setShuffleMergeAllowed(false)). Describe those active differences instead. - ShuffleManager's External-Shuffle-Service coexistence note was phrased as an incomplete, conditional sentence; reword per the reviewer's suggestion. Comment-only; no behavior change. Co-authored-by: Isaac
…ined shuffle
Native DAGScheduler support for concurrent-stage scheduling over a
PipelinedShuffleDependency (from the prior PR). A pipelined shuffle is
incrementally readable: its consumer stage may begin reading output while the
producer is still running, so the two are co-scheduled ('pipelined group')
instead of the consumer waiting for the producer to materialize.
submitStage: when a stage has missing parents, classify them by their shuffle
dependency type. A parent read through a PipelinedShuffleDependency is a
pipelined parent. The stage is co-scheduled with its producers (its tasks
submitted immediately) only if every missing parent is pipelined AND each is
already running; otherwise it parks in waitingStages exactly as before. A stage
with a regular missing parent, or a pipelined parent not yet running, waits and
is reconsidered later. This is inert for jobs with no pipelined dependency --
the full DAGSchedulerSuite is unchanged.
submitWaitingPipelinedChildStages: the 'producer started running' analog of
submitWaitingChildStages. When a pipelined producer starts, its waiting
consumers are co-scheduled immediately (not only when the producer completes),
so a consumer parked because its producer sat behind a regular shuffle is
co-scheduled as soon as the producer runs. Cascades transitively.
handleJobSubmitted: reject a job that uses a pipelined dependency when
speculation is enabled -- a speculative producer copy would race a consumer
already reading the producer's partial output, with no commit barrier. The check
runs before stage creation so a rejected job leaves no scheduler state behind.
Tests (DAGSchedulerSuite): concurrent submission, inertness for a regular
shuffle, mixed pipelined+regular parents, a deep all-pipelined chain, a pipelined
producer behind a regular shuffle (must not co-schedule early), two pipelined
parents (co-schedule and re-park semantics), fan-out reconsideration to multiple
consumers, transitive cascade, no double-submission on producer completion, and
speculation rejection (plus that a regular job under speculation is not
rejected).
Co-authored-by: Isaac
…uster capacity Best-effort admission check for pipelined groups. All member stages of a group must run concurrently, so if the group's total task demand exceeds the cluster's total concurrent-task capacity it can never be co-resident and, lacking an out-of-band slot reservation, would deadlock (the consumer holds slots waiting for producer output while the producer cannot get slots to produce). We fail fast with a clear CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT error instead. When a pipelined group is about to be co-scheduled, submitStage computes the group's full pipelined-connected component (pipelinedGroupOf) and compares its summed task demand against maxConcurrentTasksForStage (production: sc.maxNumConcurrentTasks for the stage's resource profile). If demand exceeds capacity, the job is aborted; the already-launched producers are torn down by the normal job-abort path (cancelRunningIndependentStages). This is deliberately best-effort, not the atomic gang reservation deferred to a later hardening step: it compares the whole group's demand against TOTAL capacity (not free slots), checked once at co-schedule time. That converts the common under-provisioned case (a group that can never fit) into a clear error rather than a hang; races against other concurrently admitting work are left to the future atomic version. Inert for jobs with no pipelined dependency. maxConcurrentTasksForStage is a protected seam so tests can control reported capacity without changing the cluster's core count. Tests (DAGSchedulerSuite): a group too large to co-fit fails fast with CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT; a group that fits is co-scheduled normally; producer/consumer task sets marked isPipelined; regular task sets are not. Co-authored-by: Isaac
…4.1) Upgrade the gang-admission slot check from total capacity to currently-FREE slots: free = per-profile total capacity minus tasks already running for OTHER work (other groups / regular jobs), excluding the group's own already-running members. Adds TaskSchedulerImpl.runningTasksForOtherWorkInProfile (single-lock, resource-profile-scoped, zombie-filtered) and a DAGScheduler seam. Comparing against total capacity would admit a group that fits in principle but cannot co-fit beside a busy neighbor, then hang; free-slot admission fails fast instead. Also bumps the incremental-manager config .version to 4.3.0.
… add a slot-check disable flag Two updates to the pipelined-group slot admission check (spec S4.1), per an updated spec: - Count outstanding demand, not just running tasks. The occupancy of OTHER work in a resource profile now sums each task set's not-yet-completed tasks (numTasks - tasksSuccessful), i.e. running plus enqueued, instead of only the tasks actively running. A neighbor's queued backlog is real committed demand that can starve a co-scheduled group, so charging only running tasks could admit a group that then hangs once the backlog launches. Renamed TaskSchedulerImpl's helper to outstandingTasksForOtherWorkInProfile and the DAGScheduler seam to outstandingTasksForOtherWork to reflect the new semantics. - Add spark.scheduler.pipelinedGroup.slotCheck.enabled (internal, default true). When false, pipelinedGroupExceedsCapacity is skipped entirely and a group is co-scheduled unconditionally. This mirrors ConcurrentStageDAGScheduler's ability to disable its slot check, for deployments that admit capacity out-of-band (e.g. a slot reservation) and own admission themselves. Tests: the TaskSchedulerImpl unit test now submits more tasks than slots and asserts the count includes the enqueued tasks; a new DAGScheduler test asserts that with the check disabled an over-capacity group is co-scheduled rather than failed. Co-authored-by: Isaac
… the group finishes Group-observable completion for pipelined groups (spec S5). A stage co-scheduled with a still-running pipelined producer (a pipelined consumer) must not have its successful completion processed early: doing so would advance job completion and cancel the still-running producer (via cancelRunningIndependentStages), or make the consumer's output observable before the producer's. handleTaskCompletion: near the top, before any of the event's side effects, if a Success event belongs to a pipelined consumer with unfinished pipelined producers, buffer the whole CompletionEvent and return. This defers the entire event (coarse model), so its side effects (accumulator update, TaskEnd listener event, stage/job completion) run exactly once -- at replay. markStageAsFinished: when a pipelined producer finishes, release its deferred consumers -- but only when the producer's outcome is final. A ShuffleMapStage that finished without error yet is not isAvailable (an output missing; it will be resubmitted) is NOT treated as done, so its consumers stay deferred until the reattempt makes it available. On genuine success the buffered events are replayed; on producer failure they are dropped (the group reruns, S6) but their TaskEnd events are still emitted so listeners do not leak active-task accounting. cleanupStateForJobAndIndependentStages: drop any deferral keyed on a removed stage and remove it from other consumers' pending-producer sets, so no deferral outlives its job (e.g. on abort). assertDataStructuresEmpty also checks the deferral map is empty. Inert for jobs with no pipelined dependency: the deferral map is never populated, the completion check is a always-miss map lookup, and release/cleanup are no-ops. Tests (DAGSchedulerSuite): early-finishing consumer does not end the job or cancel its running producer; normal producer-then-consumer ordering; buffered completions dropped when the producer fails; TaskEnd fired exactly once (no buffer+replay duplication); deferral released only when the producer is genuinely available. Co-authored-by: Isaac
…ts, reject dynamic allocation Two correctness fixes to the pipelined-group slot admission, found in review: - Skip zombie attempts in outstandingTasksForOtherWorkInProfile. The occupancy count summed numTasks - tasksSuccessful over every attempt in taskSetsByStageIdAndAttempt, including zombie (superseded) attempts. A retried/killed stage can have both a zombie and a live attempt in the map at once, and the live attempt already re-runs the zombie's outstanding tasks -- so counting both double-counted that stage's demand, inflating occupancy and potentially failing a pipelined group with CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT that would actually fit. Add the !isZombie filter used elsewhere on this map. - Reject a pipelined-shuffle job under dynamic allocation. A pipelined group is gang-scheduled and its free-slot check measures currently-active executors; under dynamic allocation a job can be submitted before any executor has spun up, so the group would be failed against a transient 0-slot snapshot even though the cluster would soon have capacity. Barrier scheduling forbids the same combination (checkBarrierStageWithDynamicAllocation); pipelined groups now do likewise. The former rejectSpeculationWithPipelinedShuffle is generalized to rejectUnsupportedPipelinedJob, which rejects both speculation and dynamic allocation (single RDD-graph walk, run only when a relevant feature is enabled). Tests: a zombie + live attempt is counted once, not twice; a pipelined job under dynamic allocation is rejected while a regular job under dynamic allocation is not. Co-authored-by: Isaac
…by a later PR in the stack isPipelinedGroupMember is defined here alongside the other group-topology helpers but is first used by the group-atomic failure handling added in a later PR of this stack (TaskSet.isPipelined and the member-failure fail-fast, spec S6). Reviewed in isolation this change reads as introducing an unused private method; note the forward reference in the scaladoc so it is not mistaken for dead code. Co-authored-by: Isaac
…unning plus enqueued The admission check counts OTHER work's OUTSTANDING tasks (running plus enqueued, `numTasks - tasksSuccessful`) against a resource profile's capacity -- as `outstandingTasksForOtherWork` / `TaskSchedulerImpl. outstandingTasksForOtherWorkInProfile` and the config doc already state. Two spots in DAGScheduler still described this as only the tasks "already running", which was stale after the count was widened to include enqueued tasks; the `pipelinedGroupExceedsCapacity` scaladoc even contradicted its own later "running or enqueued" line. Reword both to "outstanding -- running plus enqueued" so the prose matches the code. Doc/comment only; no behavior change. Co-authored-by: Isaac
… admit the group up front v1 (M1, real-time-mode scope) supports a job that is either all-regular or all-pipelined, not a mix. This lets an all-pipelined job's whole stage graph be one pipelined group with no regular prefix, so gang admission can be decided up front -- before any member stage is submitted -- rather than mid-DAG once a producer is already running. Changes in handleJobSubmitted (before createResultStage, so a rejection leaves no partial scheduler state, exactly like the barrier slot check and the speculation/DA reject): - classifyJobShuffleKinds walks the RDD graph and rejects a job that mixes a pipelined shuffle with a regular one, fail-fast. - rejectUnadmittablePipelinedGroup computes the whole all-pipelined group's concurrent-task demand from the RDD graph and checks it against free slots (maxNumConcurrentTasks minus other work's running-plus-enqueued demand, spec S4.1); fails the job if it cannot fit. No scheduler retry -- a transient shortfall is the caller's to retry (e.g. the streaming batch loop reruns the batch). This is true all-or-nothing gang admission: the whole group is admitted, or the job fails before any member runs, so a member is never left running while a sibling cannot get slots. The late slot check in submitStage's co-schedule branch (which measured a mid-flight snapshot after a producer was already running) is removed, along with the now-unused pipelinedGroupExceedsCapacity / pipelinedGroupOf; the capacity/occupancy seams are refactored to be resource-profile-keyed. Tests: mixed jobs rejected up front; all-PG group admitted/rejected up front (2- and 3-stage); the removed mid-DAG prefix tests are dropped as that shape is no longer supported. Co-authored-by: Isaac
…oncurrentStageDAGScheduler Rename the deferred-completion machinery to match the production RTM scheduler (ConcurrentStageDAGScheduler), so reviewers familiar with that code read the same concepts here. Pure rename; no behavior change. - pipelinedConsumerDeferrals -> dependentStageMap - DeferredCompletion -> DependentStageInfo - pendingProducers -> parents - bufferedEvents -> delayedTaskCompletionEvents The pipelined-specific helpers (isPipelinedProducer, isPipelinedGroupMember, submitWaitingPipelinedChildStages, rejectUnadmittablePipelinedGroup) keep their names -- they belong to the type-driven pipelined model and have no clean analog in the reference scheduler. Co-authored-by: Isaac
… check; fix a vacuous admission test
Two test-hygiene fixes from an adversarial review:
- assertDataStructuresEmpty now asserts dependentStageMap.isEmpty, per the
DAGScheduler class-doc checklist ("when adding a new data structure, update
DAGSchedulerSuite.assertDataStructuresEmpty ... to catch memory leaks"). The
new dependentStageMap (the pipelined deferred-completion buffer) was omitted;
it is always drained today, so this passes -- and now guards a future leak of
buffered CompletionEvents / a consumer stage left as perpetually-running.
- Retitle/reframe "the group's own running members are not charged against its
admission": it stubbed the occupancy seam to a constant ignoring
excludeStageIds, and the up-front admission passes Set.empty (no job stage
exists yet), so it never exercised member exclusion -- it only asserts a
demand==free-capacity group is admitted (a real boundary). The member-exclusion
invariant is genuinely covered in TaskSchedulerImplSuite; the title/comment now
say what the test actually verifies.
Co-authored-by: Isaac
…n, run per-task effects inline (S5.1) The group-observable completion for a pipelined consumer (spec S5) previously used a coarse model: when a consumer finished ahead of its still-running producer, the WHOLE CompletionEvent was buffered and its side effects were withheld until replay. That contradicts the negotiated spec S5.1, which requires per-task side effects -- accumulator updates and the SparkListenerTaskEnd listener event -- to flow in real time, deferring ONLY the stage/job-completion decision (advancing that early is what would cancel the still-running producer or expose the consumer's output prematurely). Switch to the fine-grained model: - handleTaskCompletion runs a deferred consumer's per-task effects inline (as for any other stage), then buffers the event and returns before the completion bookkeeping. - A new CompletionEvent.finishOnly flag marks a replayed event; on replay only the completion bookkeeping runs, so the per-task effects are never applied twice. - releaseDeferredPipelinedConsumers replays with finishOnly=true; the producer-failure drop path simply discards the buffered events (their TaskEnd already fired inline), removing the old asymmetry where the drop path re-emitted postTaskEnd but not updateAccumulators. Observability improves: a finished consumer's tasks are no longer reported as running for the producer's remaining lifetime, and no late TaskEnd burst appears at group abort. Tests: the two deferral tests now assert per-task TaskEnd fires in real time (before the producer finishes) and exactly once. Reverting the fix makes them fail (TaskEnd count 0 instead of 2 while the producer runs). Full DAGSchedulerSuite green. Co-authored-by: Isaac
…reads as unused" and a phantom symbol The NOTE on isPipelinedGroupMember said it "reads as unused when this change is viewed alone" and listed `pipelinedGroupOf` as a sibling helper. Both are inaccurate in the merged view: the method IS used (its call sites -- TaskSet.isPipelined tagging and the member-FetchFailed group abort -- are part of this stack), and `pipelinedGroupOf` does not exist. Reword to name the actual call sites and drop the phantom reference. Co-authored-by: Isaac
…private[scheduler] The method exists only to back the pipelined-group slot admission check in DAGScheduler; it does not belong on the public TaskSchedulerImpl surface. Its only callers (DAGScheduler and TaskSchedulerImplSuite) are in org.apache.spark.scheduler, so private[scheduler] is sufficient and keeps the new capacity-accounting method out of the public API. Co-authored-by: Isaac
… the retry behavior The gang-admission scaladoc likened the check to barrier's slot check. The likeness is only that both reject before any stage is created and compute demand from the RDD graph; their retry behavior differs and a reader should not infer equivalence. Spell out the contrast: barrier re-posts the job and retries its check up to a max-failures bound, whereas pipelined admission is terminal (one check, then fail) and delegates transient-shortfall retry to the caller -- scheduler-side PG-admission retry being a post-v1 refinement (spec S4.1). Co-authored-by: Isaac
jerrypeng
force-pushed
the
stack/pipelined-shuffle-pr3-scheduling
branch
from
July 20, 2026 21:04
2630bd3 to
bb0098a
Compare
…t lands on a torn-down stage A deferred pipelined-consumer completion fires its TaskEnd inline on first completion; only the stage/job bookkeeping is re-posted (finishOnly=true) once the producer finishes. The cancelled-stage guard in handleTaskCompletion (which posts TaskEnd for an event whose stage is no longer tracked) ran unconditionally and before the finishOnly check. If the group is torn down (a sibling aborts) between a finishOnly re-post and its dequeue, the replay lands on the removed stage and the guard would re-post TaskEnd -- double-counting it in listeners (AppStatusListener's active-task count could go negative). Skip postTaskEnd for a finishOnly event in that guard: its per-task TaskEnd already fired inline. A regression test on the failure PR (which has the group-atomic abort this needs) drives a finishOnly replay onto a torn-down stage and asserts no extra TaskEnd; reverting this guard makes it fail (TaskEnd count goes up by one). Co-authored-by: Isaac
…cumulator deltas are not un-merged
The drop-path comment implied full cleanliness ("only the deferred completion bookkeeping
was buffered ... must NOT be applied"). But a deferred consumer's per-task effects, including
updateAccumulators, run inline before buffering, so on a producer-failure drop the already-
applied accumulator deltas are NOT un-merged. Note that this is consistent with base Spark
(accumulators are non-transactional across abort+rerun): the rerun re-delivers under the S5
idempotent-sink contract, and RTM SQL metrics are fresh per batch, so there is no cross-batch
double-count. Comment-only.
Co-authored-by: Isaac
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changes were proposed in this pull request?
This PR adds native
DAGSchedulersupport for concurrently scheduling stages connected by aPipelinedShuffleDependency(added earlier in the stack), together with the admission andcompletion semantics that make co-scheduling correct.
A pipelined shuffle is incrementally readable: a consumer stage may begin reading the producer's
output while the producer is still running. The stock scheduler runs a consumer only after its
producer has fully materialized; this PR teaches the scheduler to co-schedule a producer and its
pipelined consumer as a pipelined group instead. Every new path is gated on a job actually
using a pipelined dependency, so behavior is unchanged for all existing jobs (the existing
DAGSchedulerSuiteis unaffected).v1 scope: a job is all-regular or all-pipelined (a job mixing the two is rejected up front).
So an all-pipelined job's whole stage graph is one pipelined group, which lets admission be decided
once, up front.
Up-front gang admission (
handleJobSubmitted). All members of a pipelined group must runconcurrently, so a group that cannot fit would deadlock (a consumer holds slots waiting for
producer output while the producer cannot get slots to produce). Before any stage is created, the
group's total task demand (computed from the RDD graph) is compared against the currently free
slots of its resource profile -- total capacity (
maxNumConcurrentTasks) minus the outstanding(running plus enqueued) task demand of other work in the same profile. Counting enqueued, not
just running, demand prevents two groups from each passing the check yet failing to co-fit. If the
group does not fit, the job fails fast with
CONCURRENT_SCHEDULER_INSUFFICIENT_SLOTbefore anymember runs -- true all-or-nothing admission that leaves no partial scheduler state (like the
barrier slot check). There is no scheduler-side retry: a transient shortfall is the caller's to
retry (a streaming query's batch loop reruns the batch). The check can be disabled with
spark.scheduler.pipelinedGroup.slotCheck.enabled=falsefor deployments that admit capacityout-of-band (e.g. a slot reservation).
Co-scheduling (
submitStage). A stage's missing parents are classified by shuffle dependencytype; a parent reached through a
PipelinedShuffleDependencyis a "pipelined parent". The stage isco-scheduled with its producers (tasks submitted immediately) only if every missing parent is
pipelined and each is already running; otherwise it parks in
waitingStagesexactly as before.submitWaitingPipelinedChildStagesis the "producer started running" analog of the existing"producer completed" hook: when a pipelined producer starts, its waiting consumers are reconsidered
immediately, cascading transitively down a chain.
Group-observable completion (
handleTaskCompletion/markStageAsFinished), fine-grainedmodel. A consumer co-scheduled with a still-running producer can finish first. Advancing its
stage/job completion early would end the job and cancel the still-running producer, or make the
consumer's output observable before the producer's. So only the stage/job-completion decision
is deferred: the consumer's per-task side effects -- accumulator updates and the
SparkListenerTaskEndlistener event -- run inline as its tasks finish, exactly as for any otherstage (per the negotiated spec's completion model), and only the completion bookkeeping is buffered.
It is replayed when the producer's outcome is final: applied on genuine producer success, dropped
on producer failure (the group reruns). Deferrals are cleaned up on job end/abort so none outlive
their job.
Other guards. A job using a pipelined dependency is rejected up front when speculation is
enabled (a speculative producer copy would race a consumer already reading partial output, with no
commit barrier) or when dynamic allocation is enabled (gang admission needs a stable slot set);
and a
PipelinedShuffleDependencycannot be submitted as a map-stage job (no durable map output tocompute statistics from).
Main changes:
DAGScheduler.scala-- up-front gang admission, co-scheduling, and completion deferral.TaskSchedulerImpl.outstandingTasksForOtherWorkInProfile-- a resource-profile-scopedoutstanding-task (running + enqueued) count for the admission check (
private[scheduler]).ActiveJob.hasPipelinedDependency-- a per-job flag so the per-submit group-membership checkshort-circuits for a regular job.
spark.scheduler.pipelinedGroup.slotCheck.enabled-- a newinternal()config (defaulttrue).CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT-- a new error condition.This is a follow-up in the stack that begins with
PipelinedShuffleDependency(#57286) and thedependency-typed shuffle-manager routing; it is the first PR that makes the scheduler behave
differently for a pipelined dependency. Group-atomic failure/rerun and additional fail-fast checks
for unsupported idioms follow in later PRs of the stack.
Known v1 limitation (admission is a check, not a reservation). Admission verifies free capacity
at submission but does not reserve slots, and member tasks launch asynchronously as offers arrive.
In the narrow window between admission and all members launching, if an executor that carries no
member task is lost (an idle one, or one running only another job), total capacity can drop with no
member-task failure to trigger a rerun, and the group could stall. This is not a data-correctness
issue (a stalled group commits nothing), it is bounded to that launch window on a warm, right-sized
cluster (dynamic allocation is rejected, so the executor set is otherwise stable), and it resolves
through the same caller-driven batch restart used for a transient admission shortfall. A
reservation-based gang launch is a post-v1 refinement (spec S4.1).
Why are the changes needed?
PipelinedShuffleDependencyand its incremental shuffle routing let a consumer read a producer'soutput as it is produced, but nothing takes advantage of that until the scheduler co-schedules the
two stages -- otherwise the consumer still waits for the producer to fully materialize and the
pipelining is never realized. Co-scheduling in turn requires admission control (a group that cannot
co-fit must fail fast, not deadlock) and completion control (a fast-finishing consumer must not end
the job or cancel its producer). This PR provides both.
Does this PR introduce any user-facing change?
No. All new behavior is gated on a job using a
PipelinedShuffleDependency, which nothing constructsyet, so for every existing job the scheduler behaves exactly as before. The new
spark.scheduler.pipelinedGroup.slotCheck.enabledconfig isinternal()and defaults totrue,and
CONCURRENT_SCHEDULER_INSUFFICIENT_SLOTis a new error condition that can only surface for a jobthat uses a pipelined dependency.
How was this patch tested?
New unit tests in
DAGSchedulerSuitecover:mixing pipelined + regular shuffles rejected up front; a deep all-pipelined chain co-scheduled;
transitive cascade when a producer starts; no double-submission on producer completion;
speculation is not rejected); a pipelined dependency submitted as a map-stage job rejected;
CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT; a group that fits total capacity but not free slotsfailing fast; a group whose demand exactly equals free capacity admitted; other work's outstanding
demand charged against admission; the slot check disabled admitting an over-capacity group;
producer; its
TaskEndfired in real time (before the producer finishes) and exactly once (nobuffer+replay duplication); normal producer-then-consumer ordering; and the deferral released only
when the producer is genuinely available.
TaskSchedulerImplSuitecoversoutstandingTasksForOtherWorkInProfilecounting running + enqueuedtasks, excluding given stages, being resource-profile-scoped, and not double-counting a
zombie + live attempt.
Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 4.8)