[WIP] [SPARK-XXXXX][CORE] Concurrently schedule pipelined-shuffle stage groups in the DAGScheduler#57341
Open
jerrypeng wants to merge 27 commits into
Open
[WIP] [SPARK-XXXXX][CORE] Concurrently schedule pipelined-shuffle stage groups in the DAGScheduler#57341jerrypeng wants to merge 27 commits into
jerrypeng wants to merge 27 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
…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
jerrypeng
force-pushed
the
stack/pipelined-shuffle-pr3-scheduling
branch
from
July 17, 2026 21:54
9aefc6d to
f7a6c5e
Compare
jerrypeng
changed the title
…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
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).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 admission / fail-fast. All members of a pipelined group must run concurrently, so a group
that cannot fit would deadlock (the consumer holds slots waiting for producer output while the
producer cannot get slots to produce). Before co-scheduling, the group's total task demand is
compared against the currently free slots of its resource profile -- total capacity minus the
outstanding (running plus enqueued) task demand of other work in the same profile, excluding the
group's own members. If demand exceeds free slots, the job fails fast with
CONCURRENT_SCHEDULER_INSUFFICIENT_SLOTinstead of hanging. This is best-effort (checked once atco-schedule time), not the atomic gang reservation deferred to a later step. It can be disabled
with
spark.scheduler.pipelinedGroup.slotCheck.enabled=falsefor deployments that admit capacityout-of-band (e.g. a slot reservation).
Group-observable completion (
handleTaskCompletion/markStageAsFinished). A consumerco-scheduled with a still-running producer can finish first. Processing that completion normally
would advance job completion and cancel the still-running producer, or make the consumer's output
observable before the producer's. So a
Successevent for a pipelined consumer with unfinishedproducers is buffered whole (coarse deferral -- side effects run exactly once, at replay) and
released only when the producer's outcome is final: replayed on genuine producer success, dropped
on producer failure (the group reruns), with
TaskEndstill emitted so listener accounting doesnot leak. 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), and a
PipelinedShuffleDependencycannot be submitted as a map-stage job (nodurable map output to compute statistics from).
Main changes:
DAGScheduler.scala-- co-scheduling, group admission, and completion deferral.TaskSchedulerImpl.scala-- a resource-profile-scoped outstanding-task count for the admissioncheck.
spark.scheduler.pipelinedGroup.slotCheck.enabled-- a newinternal()config (defaulttrue)to disable the admission check.
CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT-- a new error condition.This is a follow-up in the stack that begins with
PipelinedShuffleDependencyand 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.
Why are the changes needed?
PipelinedShuffleDependencyand its incremental shuffle routing let a consumer read a producer'ssubmitWaitingPipelinedChildStagesis 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 admission / fail-fast. All members of a pipelined group must run concurrently, so a group
that cannot fit would deadlock (the consumer holds slots waiting for producer output while the
producer cannot get slots to produce). Before co-scheduling, the group's total task demand is
compared against the currently free slots of its resource profile -- total capacity minus the
outstanding (running plus enqueued) task demand of other work in the same profile, excluding the
group's own members. If demand exceeds free slots, the job fails fast with
CONCURRENT_SCHEDULER_INSUFFICIENT_SLOTinstead of hanging. This is best-effort (checked once atco-schedule time), not the atomic gang reservation deferred to a later step. It can be disabled
with
spark.scheduler.pipelinedGroup.slotCheck.enabled=falsefor deployments that admit capacityout-of-band (e.g. a slot reservation).
Group-observable completion (
handleTaskCompletion/markStageAsFinished). A consumerco-scheduled with a still-running producer can finish first. Processing that completion normally
would advance job completion and cancel the still-running producer, or make the consumer's output
observable before the producer's. So a
Successevent for a pipelined consumer with unfinishedproducers is buffered whole (coarse deferral -- side effects run exactly once, at replay) and
released only when the producer's outcome is final: replayed on genuine producer success, dropped
on producer failure (the group reruns), with
TaskEndstill emitted so listener accounting doesnot leak. 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), and a
PipelinedShuffleDependencycannot be submitted as a map-stage job (nodurable map output to compute statistics from).
Main changes:
DAGScheduler.scala-- co-scheduling, group admission, and completion deferral.TaskSchedulerImpl.scala-- a resource-profile-scoped outstanding-task count for the admissioncheck.
spark.scheduler.pipelinedGroup.slotCheck.enabled-- a newinternal()config (defaulttrue)to disable the admission check.
CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT-- a new error condition.This is a follow-up in the stack that begins with
PipelinedShuffleDependencyand 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.
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:
and the group must complete atomically.
This PR provides all three.
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:pipelined + regular parents; a deep all-pipelined chain; a pipelined producer parked behind a
regular shuffle (must not co-schedule early); two pipelined parents (co-schedule and re-park);
fan-out reconsideration to multiple consumers; transitive cascade; no double-submission on producer
completion;
CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT; a groupthat fits total capacity but not free slots failing fast; the group's own running members not
charged against its admission; the slot check disabled admitting an over-capacity group;
producer-then-consumer ordering; buffered completions dropped when the producer fails;
TaskEndfired exactly once (no buffer+replay duplication); and deferral released only when the producer is
genuinely available.
TaskSchedulerImplSuitecoversoutstandingTasksForOtherWorkInProfilecounting running + enqueuedtasks, excluding given stages, and being resource-profile-scoped.
Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 4.8)