[SPARK-57000][CORE][SS][RTM] Add concurrent scheduling capabilites for Real-time Mode

[jerrypeng]

What changes were proposed in this pull request?

This PR introduces ConcurrentStageDAGScheduler, the scheduler needed to power real-time mode for Structured Streaming.

In real-time mode, a streaming query continuously produces output with end-to-end latency on the order of tens of milliseconds — far below the latency floor of traditional micro-batch
execution. To get there, the query has to abandon the "run stage N, materialize its shuffle output, then run stage N+1" model that the default DAGScheduler enforces. Instead, every
stage of the query must run at the same time, with records flowing from upstream tasks to downstream tasks through a streaming shuffle as they're produced.

ConcurrentStageDAGScheduler is the scheduling half of that design. Concretely, it:

• Adds a new opt-in DAGScheduler implementation, ConcurrentStageDAGScheduler, selected via spark.scheduler.dagSchedulerType=ConcurrentStageDAGScheduler and engaged per-job via the
  streaming.concurrent.stages.enabled=true property.
• Walks the stage DAG on job submission, marks all stages reachable from the final stage as concurrent, and validates that the cluster has enough free slots to run them all simultaneously
  (gated by spark.scheduler.realtimeModeSlotsCheck.disabled, with a new CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT error class for the failure path). The DAG walk accumulates into a local
  set and only commits to scheduler state after the slot check passes, so a failed submission can't leak partial state.
• Submits child stages while their parents are still running, and defers a child stage's task-completion events until every concurrent parent has actually finished — preserving the
  invariant that DAGScheduler only sees "all of a stage's parents are done" task completions, even though tasks are physically running concurrently.
• Adds the smallest extension points to DAGScheduler to make this possible: one empty onFinalStageCreated hook, two package-private accessors (isRunningStage, getStage), and
  relaxes submitStage, markStageAsFinished, submitMissingTasks, and activeJobForStage from private to protected. All hooks are no-ops for the default scheduler.
• Wires real-time-mode-aware behavior into TaskSchedulerImpl (TaskSets with concurrent stages get maxTaskFailures=1, since a streaming task failure must restart the query rather than
  silently retry against a still-running shuffle) and TaskSetManager (ExecutorLostFailure counts toward maxTaskFailures instead of being exempted, so executor loss propagates as a
  query failure rather than a silent stall).
• Both speculation guards — per-job local property and cluster-wide sc.conf.get(SPECULATION_ENABLED) — reject concurrent-stage jobs with speculation enabled, matching how the rest of
  core reads the config.
• Refactors DAGSchedulerSuite into an abstract DAGSchedulerSuiteBase + TestDAGScheduler trait so the new suite can reuse the existing scheduler test harness without duplicating it.

Why are the changes needed?

Real-time mode is the only execution model in which a Structured Streaming query can deliver sub-100ms end-to-end latency, and concurrent stage scheduling is a hard prerequisite for it.
Here's why the default scheduler can't deliver that on its own:

Sequential stage execution is the latency floor for streaming. The default DAGScheduler waits for stage N to complete — every task done, every byte written to shuffle storage, every
map output registered with the MapOutputTracker — before submitting stage N+1. For a typical streaming query with a source, a stateful operator, and a sink, that means each
micro-batch's latency is the sum of each stage's processing time plus the sum of each shuffle's serialization/deserialization cost. Even with small per-stage costs, the sum dominates
as queries get more complex, and there's no architectural way to reduce it within the existing scheduler.

Real-time mode pipes data between stages via a streaming shuffle, not a materialized one. Downstream tasks subscribe to upstream tasks' output as it's produced — there's no "stage N
is done, here are the map outputs" handoff. For that to work, all stages of the job must be running simultaneously when records start flowing. If stage N+1 isn't running yet, stage N has
no consumer for the records it produces and either drops them or blocks. So "schedule all stages concurrently" isn't an optimization for real-time mode — it's a correctness requirement of
the streaming shuffle.

Failure semantics also have to change. In batch mode, a task failure caused by an executor crash is exempted from the failure count because the executor's loss isn't the task's fault
and the framework can re-run the task elsewhere. In real-time mode that exemption is wrong: the streaming shuffle has in-flight records that can't be reconstructed, so an executor loss
must fail the query and let it restart from a checkpoint. Similarly, retrying a single task against a streaming shuffle that's already partially consumed would corrupt state — so
concurrent-stage TaskSets are capped at maxFailures=1.

The default scheduler must stay untouched for batch. Real-time mode is opt-in and additive — the cluster still needs to run batch and non-real-time streaming jobs with their existing
semantics. Hence the scheduler-type config, the per-job opt-in property, and the empty-by-default onFinalStageCreated hook on the base DAGScheduler: when nothing opts in, nothing
changes.

Does this PR introduce any user-facing change?

No user-facing behavior change for any existing workload. Without setting the new config, SparkContext builds the same DAGScheduler it always has, and the default scheduler's
behavior is unchanged.

The PR does introduce two new internal configs (both internal(), so not part of the public surface):

• spark.scheduler.dagSchedulerType — chooses the DAGScheduler implementation. Defaults to "DAGScheduler".
• spark.scheduler.realtimeModeSlotsCheck.disabled — skips the slot-availability check used by the concurrent scheduler. Defaults to false.

And one new error class:

• CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT — thrown by the concurrent scheduler when a streaming job needs more concurrent slots than the cluster offers.

How was this patch tested?

Added one new test suite plus targeted regression tests in the existing TaskScheduler suites:

1. ConcurrentStageDAGSchedulerSuite — exercises the new scheduler end-to-end through the existing DAGSchedulerSuiteBase test harness. Tests cover:

   • Happy paths. Simple two-stage concurrent job (both stages enter runningStages on submission; the child's task completions are buffered until the parent finishes). Complex
     six-stage DAG with diamond dependencies (verifies parent-tracking, deferred-event buffering, and correct release order when parents finish out of order). Concurrent stages disabled in
     properties (scheduler falls back to default sequential behavior).
   • Failure-path cleanup. Slot-check failure leaves internal state empty (no leak of partially-visited stages). Stage abort with a shared parent stage cleans up dependentStageMap
     (the parent isn't marked finished by the cascade, so the cleanup at the end of markStageAsFinished is the only path that can release the entry). Job cancellation cleans up both
     concurrentStages and dependentStageMap. Executor-loss-induced abort (via the maxFailures=1 path) cleans up both maps.
   • Speculation rejection in two variants: per-job local property and cluster-wide SparkConf. Both cause the job to fail on submission with a clear error.
   • An extraEmptyChecks hook on assertDataStructuresEmpty is overridden to assert both concurrentStages and dependentStageMap are empty, and is called from afterEach, so
     every locally-defined test and every inherited test from DAGSchedulerSuiteBase (149 of them) automatically validates the new state invariants. Total: 155 tests pass in this suite.

   By inheriting from DAGSchedulerSuiteBase, the suite also runs all 149 existing DAGScheduler tests against ConcurrentStageDAGScheduler — free regression coverage that the new
   scheduler behaves identically to DAGScheduler when concurrent mode is not engaged.

2. TaskSchedulerImplSuite — one new test: a TaskSet with streaming.concurrent.stages.enabled=true is submitted with maxTaskFailures=1 regardless of spark.task.maxFailures; a
   regular TaskSet still gets the cluster default. Regression-guards both branches of the new conditional.

3. TaskSetManagerSuite — two new tests covering the new failure-counting behavior:

   • With concurrent stages enabled, an ExecutorLostFailure with exitCausedByApp=false counts toward maxTaskFailures (the query restarts rather than silently absorbing executor
     loss).
   • Without concurrent stages, the same failure does not count — regression guard for the default behavior.

Full run: core/testOnly *DAGSchedulerSuite *ConcurrentStageDAGSchedulerSuite *TaskSetManagerSuite *TaskSchedulerImplSuite → 489 tests, all pass.

Was this patch authored or co-authored using generative AI tooling?

Co-authored with Claude Code (Claude Opus 4.7)

[jiangxb1987]

When this throws, the stages added to concurrentStages above are leaked — handleJobSubmitted catches the exception and fails the job, but nothing ever clears those entries. A subsequent job whose stages share IDs (e.g. retries from the same RDDChain) would inherit them. Either clear concurrentStages of the stages just visited before throwing, or capture them in a local set and only commit to concurrentStages once the slot check passes.

[jerrypeng]

fixed by accumulating into a local visitedStages set during the DAG walk and only committing to concurrentStages after the slot check passes

[jerrypeng]

Though the actual affect of this would likely be small since this would only occur on query failure.

[jiangxb1987]

The dependentStageMap cleanup path only fires when a stage in the map is named as a parent via markStageAsFinished(parent). If a dependent stage itself aborts mid-job (e.g. its single allowed failure under maxTaskFailures=1), its own entry — including any buffered delayedTaskCompletionEvents — is never removed from dependentStageMap. With concurrent jobs sharing a long-lived scheduler instance, that's a slow leak across queries. Consider clearing the entry for stage itself inside markStageAsFinished (especially when errorMessage.isDefined).

[jerrypeng]

I will just remove the stage's own entry at the end of markStageAsFinished

[jiangxb1987]

QUERY_ID already exists and is what StructuredStreamingIdAwareSchedulerLogging uses to log streaming query IDs. Adding STREAMING_QUERY_ID creates a parallel key for the same concept. Suggest dropping this addition and using LogKeys.QUERY_ID at all the callsites, or update the callsites in StructuredStreamingIdAwareSchedulerLogging.

[jerrypeng]

A query id and streaming query id are typically not the same. query id for a batch query is simply a transient id for a batch. The streaming query id is persistent for the entirety of the streaming query execution.

I would keep it here and fix it in StructuredStreamingIdAwareSchedulerLogging

[jerrypeng]

@jiangxb1987 thank you for the review! I have addressed your comments and added additional tests. PTAL.

[jiangxb1987]

This should be IllegalStateException instead of RuntimeException?

[mridulm]

Why is this not modelled using barrier scheduling ?
IIRC I had raised this question in the SPIP as well

(Marking as requested change to prevent accidental merge)

[jerrypeng]

Thanks @mridulm for the question, and for your interest in this work!

The short answer is that barrier execution mode and concurrent stage scheduling solve orthogonal problems. As I understand it, barrier mode is gang scheduling for the tasks within a single stage: it launches all N tasks of that stage simultaneously, and the tasks can then coordinate with each other mid-execution via barrier() / allGather() (MPI-style).

What real-time mode needs is different — the ability to schedule multiple stages of a job to run concurrently (which is what this PR focuses on), so records can stream from upstream stages to downstream stages through a streaming shuffle. There's no hard requirement for all tasks to coordinate, or to be co-scheduled, before the query starts.

Your question — whether RTM could benefit from gang scheduling — is a fair one. I think the answer is "maybe, but not strictly necessary." The streaming shuffle implements a backpressure mechanism that serves a similar purpose: if a downstream consumer isn't ready yet, the upstream producer backs off rather than failing, thus a coordinate execution system like barrier scheduling is not needed

[mridulm]

From a scheduler perspective, it appears that the primitives required to make this work already exists - whether it gets leveraged for mpi or real time mode is an implementation detail.

Having said that, if we are robust to not requiring all stages/tasks to be running before execution (which was my previous understanding as per proposal - perhaps I misunderstood) - what is the gap w.r.t launching a bunch of long running map stages ?

Essentially, I am trying to understand why this needs to be in the scheduler - versus in integration code

To put it differently, I am trying to make sure we have robustly considered alternatives - and I want to understand their tradeoffs; they are not clear to me

[jerrypeng]

@mridulm — This change is not needed when the streaming query is a single stage: a single long-running (map) stage runs fine on the existing scheduler, which is exactly why RTM support for single-stage stateless queries already shipped in 4.1 — no scheduler change required there.

However, multi-stage queries (e.g. stateful queries) are today executed one stage at a time, with each stage's shuffle output fully materialized before the next stage starts. To reach millisecond-level latencies, we instead need the stages of a single query to run concurrently, connected by the streaming shuffle (currently being merged in incrementally). Enabling concurrent execution of dependent stages within a single job is what requires the scheduler change — which is why this work is needed.

[mridulm]

@jerrypeng you can launch N number of map stages - and wire them to talk to each other for multi stage queries ?
Whether concurrent execution is possible depends on resource availability.

[jerrypeng]

@mridulm we can do that, but that is probably not the most eloquent or simplest solution.

A multi-stage query isn't a list of independent map stages — it's a connected stage DAG the planner already produces: shuffle-map stages at each exchange, a result stage at the sink, real shuffle-dependency edges between them, and branching where it exists (e.g. a join stage reads two shuffle inputs). To emulate that as "N map stages wired together" we'd have to re-cut the plan into separate jobs, pre-mint and inject shuffleIds to recreate the dependency edges, and re-implement cross-job failure/cancellation/completion so the micro-batch still behaves as one unit — i.e. re-derive the DAGScheduler's own stage decomposition and job coordination in streaming code, to fake a DAG we already have natively.

So instead of faking it, we keep the real plan and its real stages, and change the only thing that's actually different in real-time mode: when the stages run. Normally a stage waits for its parent to finish; here the stages run at the same time, connected by the streaming shuffle. That's a scheduling decision about an existing DAG, which is why it belongs in the scheduler.

[mridulm]

@jerrypeng Everything you describe cam be entirely in the integration layer - instead of at the scheduler.
Note - I am not necessarily against the idea of doing this change : but the semantics of what is being proposed is not clear, or fleshed out IMO.

We are letting the implementation details for a usecase define what the scheduler should look like, and making surgical changes to adapt to it - instead of defining what the semantics need to be (perhaps this has been done : but it is unclear from this PR anyway).

[jerrypeng]

@mridulm

Everything you describe cam be entirely in the integration layer - instead of at the scheduler

Can you elaborate?

but the semantics of what is being proposed is not clear, or fleshed out IMO

Sorry to hear that! Can you help me understand better what is not clear to you?

[mridulm]

Can you elaborate?

For example, this integration could be modeled as I described above - submitting map stages with streaming shuffle wired up between the stages.
(I am not arguing it is ideal - but most likely possible)

Sorry to hear that! Can you help me understand better what is not clear to you?

The proposed scheduler is codifying expectations specific to this implementation - and not generic constructs.
For example, in the default DAGScheduler, there is a rationale for why a child waits for parent before it starts (this is different from MR, which has its own model)
Here, it is unclear why/.when it can start, how 'deep' can it start, etc - the decisions appear to be driven by RTF implementation details, and not robustly defined - making it less extensible to use for other usecases. As an example, barrier scheduling was designed for ML apps, but the constructs in scheduler are generic, and applicable to other usecases as well.

In other words, I want to make sure we make scheduler changes only when required, where the behavior is not implementation details in service of an initiative, and constructs help unlock a larger class of usecases.

[jerrypeng]

@mridulm thank you for clarifying! Some background may help.

Real-time mode is a new execution mode we're introducing in Structured Streaming that lets streaming queries process data with end-to-end latencies in the milliseconds. Reaching that requires a few changes to how queries execute; scheduling is one of them, and that's what this PR covers. For a query with multiple stages to hit millisecond latencies, the cluster has to run the tasks of all stages at the same time, with adjacent stages connected by a streaming shuffle (implemented in a separate set of PRs). That lets data flow continuously through the query DAG instead of one stage at a time — and processing one stage at a time is a core reason the current model can't reach these latencies.

On the changes in this PR: the change to the existing DAG scheduler is small and additive — a no-op hook, a couple of accessors, and a few visibility relaxations — and the default behavior is unchanged. The new capability lives in a separate class (ConcurrentStageDAGScheduler), so the real-time-mode logic is isolated from the shared scheduler.

On the semantics you asked about: the new capability is that the stages of a DAG can run concurrently rather than one at a time. Today the DAG scheduler treats a data dependency between two stages as a reason to run them sequentially — the child waits for the parent. But that sequencing isn't inherent to the dependency; it's a consequence of the shuffle being materialized, where the consumer can't read anything until the producer has written its complete output. With a streaming shuffle that the consumer reads incrementally, the producer and consumer stages can run at the same time. So the semantic this introduces is: a directional data dependency constrains the ordering of data, not necessarily the concurrency of execution — and when the connecting shuffle supports incremental reads, dependent stages may execute concurrently. I hope that clarifies it.

This is a general scheduling capability — concurrent execution of data-dependent stages over an incrementally-readable shuffle — and real-time mode is simply its first consumer.

I think the question next should be how to natively integrate this into the DAGScheduler so users don't need to specify to configure to use the ConcurrentDAGscheduler to get this capability. That is what I will be working on next.

[jerrypeng]

@mridulm do you have any additional concerns I can address?

[mridulm]

I am not in favor of merging this PR.
This comment is directionally better aligned with how we should approach it.

Strawman proposal - extend support for realtime shuffle as a first class concept within DAGScheduler.

Currently we have:

• Narrow dependency between RDD's - merge into the same stage.
• Shuffle dependency - introduce a shuffle split (unless it is provable that we can convert it to narrow dependency).

With semantics around how to handle failures, etc.

Extend this to support real time shuffle as a first class support, and define :
a) Given a job, how it gets 'split' into stages and wire them based on real time shuffle dependency (when to split, when to combine within stage ? How it interacts with everything else ? For unsupported idioms - explicitly fail fast)
b) Which stages can be concurrently executed and which need to wait - define semantics around progress, starvation, etc.
c) What are the semantics around failures
d) How does this interact with existing constructs (for ex: if there is 'regular' shuffle dependency ? throw exception ? supported ?)

A lot of these are already in the current PR - we need to simply formalize them, and integrate into existing machinary.
This PR is good to test things out and validate ideas - but not for merging into Apache Spark itself

[jerrypeng]

@mridulm thank you for the detailed feedback — I think we're aligned on the destination, and I'd like to propose reaching it incrementally.

I agree the end state is: these scheduling semantics supported by the default DAGScheduler, with richer, more fine-grained abstractions — e.g. annotating in the query plan which shuffles can be read incrementally, rather than an opt-in flag. My question is whether we can sequence it into milestones rather than land it all at once.

IMO this PR already declares clear semantics for the new scheduling capability, and they're fairly generic:

1. The shuffle connecting two concurrent stages is read incrementally: the consumer reads from a still-running producer instead of waiting for fully materialized output.
2. Because of that, stages with a data dependency can run concurrently rather than sequentially.
3. Because that incremental shuffle is transient (its data can't be replayed), any task failure restarts the whole job.

None of these reference streaming — real-time mode is just the first caller, and the capability isn't streaming-specific: any feature that uses an incrementally-readable shuffle can opt into the same semantics. The PR gates them behind a streaming-named property for expedience (streaming.concurrent.stages.enabled) — happy to rename it to something more generic if you'd like.

I'd also note the DAGScheduler footprint is deliberately small: the base-class change is a no-op hook, a couple of accessors, and a few visibility relaxations, with the default execution path unchanged — precisely because I share your concern that changes there are high-risk. The new behavior is fully opt-in, so structuring it this way keeps the blast radius small: unrelated queries can't be affected by these changes. That's also what makes landing it incrementally low-risk.

Could we use this PR as the first milestone and merge it as-is? It would let us test and validate real-time mode end-to-end in-tree while we design the deeper integration. As an immediate follow-up, I will work through how to make these semantics more natively defined in the DAGScheduler.

What do you think?

[jerrypeng]

@mridulm To give a sense of what I'm picturing for the more native design:
Entry point — a capability flag on the shuffle. Add a field to ShuffleExchangeExec:

  // true if this shuffle can be read incrementally
  // (the consumer can read from a still-running producer)
  incrementalHint: Boolean

It's set during physical planning (this is an execution concern, not a logical-plan one), flows into the ShuffleDependency the exchange creates, and is read by the DAGScheduler at stage-creation time. That single flag is the entry point that opts a shuffle edge into concurrent scheduling over an incremental shuffle (e.g. the "streaming shuffle" we're building for RTM).

How RTM sets it. A physical-planning rule for RTM streaming queries marks the query's shuffle exchanges with incrementalHint = true. Nothing here is streaming-specific: any feature can write its own physical-planning rule to opt a job into concurrent scheduling + incremental shuffles, so the capability is generic — RTM is just the first caller.

Semantics implied by an incremental shuffle:

1. The stages on either side of the shuffle can run concurrently — the consumer reads from a still-running producer instead of waiting for fully materialized output.

2. Because the incremental shuffle's data is transient (it can't be replayed), a task failure in either of the co-scheduled stages reruns them (for RTM, the query restarts from checkpoint).

A cleaner generalization — separate "incremental" from "persistent". Semantic (2) is really a consequence of a second, orthogonal property: whether the shuffle data is durable/replayable. We can make that explicit with a second flag:

  // true if the shuffle data is durable/replayable, so a single task
  // failure does not require rerunning the producer stage
  persistentHint: Boolean   // name TBD

That decouples two concerns:

• incrementally-readable → enables concurrent scheduling.
• persistent/replayable → determines failure recovery: a transient shuffle means rerun the co-scheduled stages; a durable shuffle allows normal fine-grained recovery (re-read from the persisted data / offset).

Though I would defer implementing the "persistentHint" capability until there is an actual use case / implementation.

The streaming shuffle is incremental and transient. But a Kafka-backed shuffle, say, could be incremental and persistent — concurrent scheduling without the rerun-everything-on-failure penalty, since the consumer can replay from an offset. Splitting the two flags lets the construct compose across those cases instead of hardcoding RTM's failure model.

Pluggable shuffle implementation via config. The engine maps the capability to a concrete ShuffleManager. Today we dispatch between sort and streaming managers on a per-job property; the cleaner version selects per-dependency from incrementalHint, with the implementations configured:

  spark.shuffle.manager             = org.apache.spark.shuffle.sort.SortShuffleManager
  spark.shuffle.manager.incremental = org.apache.spark.shuffle.streaming.StreamingShuffleManager

A shuffle with incrementalHint = true is served by the configured incremental manager; everything else by the default — keeping the scheduler construct generic while the shuffle implementation stays pluggable.

Let me know what you think though I would still prefer to do this incrementally like my previous comment.

[mridulm]

Let me rephrase it - why do we want to add this workaround ? And not integrate it into DAGScheduler ?
Are we concerned about stability ? Incompleteness of implementation ? Side effects/interaction concerns ?

Any scheduler change suffers from potential risks - and I am trying to reason about why we are special casing here and introducing an entirely new dag scheduler.

If we want to make dag scheduler pluggable - that is a design in itself, and needs to be thought through - not as a derisking mechanism for a specific feature.

[jerrypeng]

@mridulm the implementation presented in this PR is not really a workaround — it is a working solution for the needs of RTM.

I agree that we can eventually design a more natively integrated solution that provides more generic functionality. However, can we approach that incrementally?

My philosophy for software development is iterative. I would like to first introduce something that works for the RTM use case, while minimizing risk to existing Spark use cases. That is what this PR is trying to do. The changes are intentionally scoped so that we can test RTM end-to-end without requiring a larger DAGScheduler redesign upfront.

I would rather get something working first, validate it end-to-end, and then iteratively refine the abstractions. Building a more generic framework may take time, and I am happy to work toward that, but I do not think we should block RTM progress on having the fully generalized design in place from day one.

Let me know what you think. Regardless I am going to look into how to natively built this functionality in the DAGScheduler.

[mridulm]

As I said before, I would suggest working towards integrating into DAGScheduler with a principled solution; rather than trying to introduce extension mechanisms just to derisk the work : that is a feasible approach for a prototype.

If we want to make DAGScheduler pluggable - that could be a design in itself - and should be approached as such. To be honest, there is a case to be made for that given the complexity.

I am -1 on this specific direction for the change until then.

+CC @jiangxb1987 as you reviewed this PR as well.

[jerrypeng]

@mridulm what do you think about the approach I described here:

#56055 (comment)

[cloud-fan]

Joining late — I've read the thread and I'm with @mridulm on direction: this should be a principled construct inside DAGScheduler, not a separate scheduler picked by config (pluggable DAGScheduler is its own design problem, not a derisking mechanism for one feature). But I think @jerrypeng's latest sketch (comment) is most of the way there, and this PR is a useful prototype for nailing down the semantics. Let me try to converge them.

The new ability is cross-stage gang scheduling, and that already implies a streaming shuffle. Co-scheduling two data-dependent stages is only useful if the edge is readable before the producer finishes — i.e. a streaming shuffle. So "run these stages concurrently" and "the shuffle is incremental" aren't orthogonal; they're one decision seen from two sides. The incrementalHint/persistentHint split models a four-state space where only one cell is real today (incremental + transient = the streaming shuffle), and persistentHint has no implementation behind it — I'd drop it until a persistent-incremental shuffle (e.g. Kafka-backed) actually needs it.

One nuance on "barrier": what's barrier-like is the resource side — the co-scheduled stages must all get slots at once or fail fast, which is exactly what this PR's CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT reimplements on top of the existing checkBarrierStageWithNumSlots. What's not barrier-like is execution — backpressure handles readiness, no sync point. So it's gang-scheduled stages + pipelined execution. (I'd avoid the word "barrier" in the API since it already means intra-stage MPI scheduling, but the lineage — generalizing barrier's gang-resource model from tasks to a stage set — is worth stating.)

Concretely, I'd express @mridulm's (a)–(d) as one first-class dependency kind rather than two hint flags: a pipelined shuffle dependency, peer to narrow/shuffle deps, set in physical planning and carried into the ShuffleDependency.

Add: the marker on ShuffleDependency; treat a pipelined edge as non-sequencing in stage creation so the connected component forms one co-scheduled group (a); submit the group together through a generalized barrier slot check (b); a group-failure policy where any failure reruns the group (c); fail-fast where a regular shuffle dependency crosses the group boundary (d).

Remove: ConcurrentStageDAGScheduler, spark.scheduler.dagSchedulerType, and the per-job streaming.concurrent.stages.enabled property (replaced by the marker); CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT + realtimeModeSlotsCheck.disabled (subsumed by the slot check); and re-key the TaskSchedulerImpl/TaskSetManager failure changes on group membership rather than the streaming property. RTM then becomes a pure consumer — a physical-planning rule marks its exchanges pipelined, no streaming-specific code in the scheduler.

On out-of-order completion: with stages concurrent, the result stage's tasks can finish while parents run, and the base Success path would declare the job done and cancel the still-running parents. The current PR buffers a stage's completion events until its parents finish, then replays them. I'd keep that approach — the all-or-nothing group-failure model means we never need partial-progress correctness inside a group (the only thing a native rewrite would buy), and the buffer keeps the batch-critical completion paths untouched while naturally holding a finished-but-still-pipelining stage in runningStages. Just make the group the first-class concept — an explicit group-completion contract implemented via the deferral — so it's a documented mechanism, not a free-floating shim. Worth a test for the replay window (a failure between parent-finish and replay → group reruns, buffer dropped).

Sequencing: this still gets us there incrementally, just cut differently — land the dependency type + group scheduling + group failure as the generic milestone with a non-streaming test, then add the RTM rule + streaming shuffle on top. The scheduler primitive is the hard-to-revise part, so I'd rather get its shape right in-tree first with RTM as the validating consumer than merge the subclass and re-cut it later. Happy to help review the dependency-type design.

[mridulm]

"The new ability is cross-stage gang scheduling, and that already implies a streaming shuffle"

This is a fair extension to build - and if this is the modeling exercise - we should be looking to extend the existing gang scheduling construct. I started off asking about this here :)

flows into the ShuffleDependency the exchange creates, and is read by the DAGScheduler at stage-creation time.

This is not a hint from scheduler perspective.
We should make it explicit and call it StreamingShuffleDependency or some such - and define DAGScheduler contract on how it handles the DAG

Post correctness validation (supported modes of wiring DAG) this would determine how stages are eligible for schedule + how they get scheduled (@cloud-fan 's articulation of cross stage gang scheduling for example, or something similar perhaps), straggler handling, failure handling, etc.

Concretely defining the model, behavior is my main ask - we might not build for all corner cases in first ver, and that is perfectly fine - as long as we explicitly fast fail on unsupported behaviors (Wenchen gave some good examples of this)

[viirya]

Thanks for the detailed prototype and the design discussion in the thread. I've read through it and I'm aligned with @mridulm and @cloud-fan on direction: the co-scheduling capability should be a first-class construct inside DAGScheduler, not a separate scheduler selected by spark.scheduler.dagSchedulerType. Making DAGScheduler pluggable is its own design problem and shouldn't be introduced as a derisking mechanism for one feature. @cloud-fan's convergence sketch — a first-class pipelined shuffle dependency that forms a co-scheduled group, with the group as the unit of slot-checking and failure — looks like the right shape to me, and it subsumes most of the machinery this PR adds by hand. I'd rather get that primitive's shape right in-tree with RTM as the validating consumer than merge the subclass and re-cut it later.

A few things I'd want addressed regardless of which form this lands in, since they're real today:

The maxFailures=1 justification doesn't cover the FetchFailed path. The comment at the end of markStageAsFinished says stage retries aren't supported because TaskSchedulerImpl pins maxFailures=1. But maxFailures only governs task-level retries — the FetchFailed -> resubmitFailedStages -> submitStage path in the base DAGScheduler is a separate mechanism and isn't overridden here. Since markStageAsFinished unconditionally does dependentStageMap.remove(stage) (including on the willRetry = true path), if a concurrent stage ever hits that path its buffered completion events and parent-tracking are dropped. In RTM this may be unreachable depending on the streaming shuffle, but the comment overstates the guarantee. Please either narrow it to "task-level retries" and state why FetchFailed can't occur (or should fail-fast) for a pipelined edge, or handle it explicitly. This is exactly the "interaction with a regular shuffle dependency" question @mridulm raised — the group-boundary fail-fast in the convergence sketch would answer it cleanly.

The activation gating is asymmetric. onFinalStageCreated engages concurrent scheduling only when queryBatchId.nonEmpty && isConcurrentStagesEnabled (needs queryId + batchId + the flag), but TaskSchedulerImpl (maxFailures=1) and TaskSetManager (ExecutorLostFailure counting) key off isConcurrentStagesEnabled alone. A job with the flag set but no streaming query/batch id gets the stricter task-failure semantics without concurrent scheduling — a half-activated state. Re-keying the task-scheduler changes on group membership (per the sketch) removes this; short of that, the three sites should share one predicate.

A replay-window test is worth adding. With the deferral mechanism, there's a window between a parent finishing and its child's buffered events being replayed. A failure landing in that window drops the buffer and should rerun the group. That path isn't covered today and is the subtlest part of the mechanism — @cloud-fan flagged the same thing.

The slot check compares against defaultParallelism(), which isn't slots. CoarseGrainedSchedulerBackend.defaultParallelism() returns spark.default.parallelism when set, otherwise total core count — so with that config set, CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT is checked against an unrelated value. If this survives as its own check rather than folding into a generalized checkBarrierStageWithNumSlots (which the sketch proposes), it should be documented/guarded. The barrier slot check is the existing precedent worth reusing here.

Minor:

• .version("4.2.0") on both new configs is stale — branch-4.x is at 4.3.0 now, so a normally-backported change should be 4.3.0.
• totalNumCoreForStage's match on Stage has no default case (Stage isn't sealed) — a MatchError waiting to happen if a third subclass ever appears.
• The new tests use the raw property literals "sql.streaming.queryId" / "streaming.sql.batchId" rather than the StructuredStreamingIdAwareSchedulerLogging constants the main code was updated to use.

[jerrypeng]

Pipelined Shuffle Dependency & Concurrent Stage Scheduling

A spec for running data-dependent stages of a single job concurrently, connected by a shuffle the
consumer reads incrementally.

---

1. Motivation

Today a multi-stage job runs one stage at a time: each shuffle is fully materialized before the next
stage starts. Some workloads need the stages of a single job to run concurrently, connected by a
shuffle whose consumer reads the producer's output as it is produced rather than after the
producer finishes. This spec introduces the scheduler primitives to express and run that.

"Run these stages concurrently" and "the connecting shuffle is incremental" are the same decision
seen from two sides: co-scheduling a producer and consumer is only useful if the edge is readable
before the producer completes.

---

2. Primitives

2.1 Pipelined shuffle dependency (PSD)

A shuffle dependency declared incrementally readable: a consumer stage may begin reading its
output while the producer stage is still running.

• It is a shuffle dependency (has a shuffleId, partitioner, map/reduce sides); the pipelined
  property is a binding part of the scheduler contract, not an advisory hint.
• The property is set during physical planning (an execution concern, not a logical-plan one)
  and carried into the ShuffleDependency the DAGScheduler reads at stage-creation time.

A regular shuffle dependency (RSD) is an ordinary shuffle dependency: its output must be fully
materialized before any consumer reads it.

2.2 Pipelined group (G)

The set of stages connected to one another through pipelined edges — the connected component of the
stage DAG when only pipelined edges are considered.

• A stage with no incident pipelined edge is a singleton group and behaves exactly as a normal
  stage today.
• The group — not the edge or the individual stage — is the unit of admission, slot
  checking, completion, and failure.

External input of G: a regular shuffle dependency whose consumer is in G and whose producer is
not — i.e. a normal materialized parent of the group.

---

3. Group formation

• Stage decomposition is unchanged. A pipelined dependency introduces a shuffle boundary exactly
  as a regular one does; the set of stages and their partitioning are identical. The pipelined
  property changes only when stages run relative to one another, never how the plan is cut into
  stages.
• Group = connected component over pipelined edges. As stages are created, two stages joined by a
  pipelined edge are placed in the same group; the group is the transitive closure.
• Every stage belongs to exactly one group (singletons included). Group membership is fixed at
  stage-creation time.

---

4. Scheduling & admission

• A pipelined edge is non-sequencing. The consumer of a pipelined dependency does not wait for
  its producer to materialize. (A regular-dependency consumer still waits — the default behavior.)
• Group readiness. A group is ready to be admitted when every external input of the group (its
  regular materialized parents) is available — the same precondition a normal stage has today, lifted
  to the group. Pipelined parents inside the group impose no readiness precondition.
• Gang admission (all-or-nothing). A group is admitted only if the cluster can currently run all
  tasks of all member stages concurrently; admission then submits every member stage at once.
  There is no partial admission — a group is never left with some members running while others wait
  on slots the running members occupy.
• Slot check. The group's aggregate concurrent-task demand — the sum of numTasks over member
  stages — is compared against the number of available slots in the cluster (a slot is one task's
  worth of capacity, so this is the maximum number of tasks that can run at once). If the group
  needs more slots than the cluster can offer, the submission fails fast, since the group could
  never become fully co-resident.
• Co-residency. Once admitted, all member stages of a group are simultaneously running.
• Single ownership. A pipelined group belongs to exactly one job; a pipelined producer is not
  shared across jobs.

---

5. Completion

• Group-observable completion. A member stage is not job-observably finished until all member
  stages of its group have completed successfully. A result stage's job therefore completes only when
  its entire group has, and a member that finishes ahead of the others cannot advance job or stage
  completion until the group does.
• Defer the finish decision, not per-task work. When a member finishes before its group, only its
  stage-finish / job-finish transition is deferred until group completion. Per-task side effects that
  must always run — accumulator updates, output-commit coordination, task-end listener events — run
  immediately.
• Replay window. There is a window between a member finishing and group completion. A failure in
  that window is a group failure (§6): the deferred finish transitions are dropped and the group
  reruns.

5.1 Observable completion events (listener bus)

The listener bus is an external contract that monitoring tools depend on, so it is worth stating
exactly when each event is delivered. The rule follows directly from atomic commit: task-level
events flow in real time, but stage-completion and job-completion events are held until the group
commits — so a listener never observes a member as successfully completed before the group as a
whole has.

Event	Timing	Rationale
SparkListenerTaskStart / SparkListenerTaskEnd	Real time, as they occur	Per-task facts are true when they happen; a group's members genuinely run concurrently. Deferring these would freeze a member's live progress and metrics for the whole group's duration. Note a successful TaskEnd means "this task finished," not "its output is committed" — already true in Spark, since a stage attempt can later be discarded.
SparkListenerStageSubmitted	Real time, at group admission	All member stages are submitted together (§4); a monitor should show them active simultaneously.
SparkListenerStageCompleted	Deferred to group commit; on group failure, emitted with a failure reason	"Completed" should track commit, which is atomic at the group level. A member whose tasks finish early is reported as still running until the group commits — which matches the truth that its results are not usable until then. This avoids emitting a success-shaped completion for an attempt that a later group failure would discard.
SparkListenerJobEnd(JobSucceeded)	At group commit only	Job completion delivers results to the caller and cancels sibling stages; emitting it before the group commits risks double/inconsistent result delivery if the group later fails. Non-negotiable.
SparkListenerJobEnd(JobFailed)	On group failure	Group-atomic failure (§6): buffered success transitions are dropped, never replayed as success.

Failure-path consequence: because stage- and job-completion are deferred (never emitted early as
success), a group failure needs only to emit them as failure — there is no premature success event
to retract. Already-emitted task events stand as-is (they were true), consistent with how Spark
treats task events from a stage attempt that is later discarded.

---

6. Failure

• Failure is group-atomic. Any task failure in any member stage — including a fetch failure —
  fails the whole group. There is no independent single-stage failure within a group. Two mechanisms
  make this hold:
  • No single-stage resubmit inside a group. The ordinary "resubmit one failed stage" path is
    invalid for a group member: the transient pipelined shuffle cannot be re-read, and members are
    co-scheduled. A member failure routes to group failure instead.
  • Teardown is by group membership, not producer availability. At the end of a batch the producer
    finishes and registers all its map outputs — becoming "available" — while the consumer is still
    draining the remainder, so a producer-available-while-consumer-still-running window is normal,
    not rare. For a transient pipelined shuffle, "available" does not mean the consumer is safe: the
    data is not durably stored and the consumer is still actively reading it, so a producer failure
    in that window still strands the consumer. Failure propagation that keys off producer
    availability would miss the consumer here; teardown is therefore keyed on group membership, so a
    producer failure always tears down its co-scheduled consumers regardless of the producer's
    availability at the failure instant.
• Recovery. Group failure tears down all member stages atomically and discards their deferred
  completions. Because the shuffle is transient it cannot be partially recovered; the job fails and
  the caller re-runs it.

---

7. Interaction with regular dependencies

• Regular shuffle into a group — supported. A regular dependency from outside the group to a
  member is a normal materialized parent (an external input); the group waits for it.
• Regular shuffle out of a group — supported. A regular dependency from a member to a stage
  outside the group produces a materialized output that a downstream group consumes with normal
  sequencing; the downstream waits for the group to complete.
• Regular shuffle internal to a group — fail-fast / unsupported. If a regular dependency's
  producer and consumer both fall in the same group, the group would co-schedule them concurrently
  while the regular edge demands the producer be materialized first — a contradiction. Reject at
  stage creation.
• Pipelined edges are intra-group by construction, so they never cross a group boundary.

---

8. Activation

• Activation is by group membership. Every behavioral change — admission, completion, and
  failure handling — keys off a single question: is this stage a member of a pipelined group? There
  is no separate enable flag and no half-activated state.
• Opt-in is expressed by the presence of a pipelined dependency, which a physical-planning rule
  sets on the relevant exchanges. The scheduler primitive is generic; any feature can write such a
  rule to opt a job into concurrent stage scheduling over an incremental shuffle.

---

9. Fail-fast on unsupported idioms

A pipelined group that involves any of the following is rejected at stage/group creation with a clear
error, rather than silently mis-scheduled. Each is a documented limitation of the first version.

The rejected idioms fall into two kinds. Some are moot under the group failure model (§6):
Spark's stage-level recompute/rollback mechanisms never fire inside a group, because any failure
aborts the whole group and the caller reruns from scratch — so a mechanism whose only job is to
recompute or roll back a stage after a partial failure is never reached. We reject these rather than
carry dead, self-contradictory machinery. The rest are incompatible with concurrent execution
itself and would corrupt or hang a group that never fails.

Rejected condition	Kind	Why rejected
Barrier execution in a member stage	incompatible	Barrier exposes output only after a global sync, contradicting concurrent partial reads.
Dynamic resource allocation	incompatible	Gang admission needs a stable slot set; reclaiming executors from a pinned-open group can deadlock it.
Speculative execution	incompatible	A speculative producer copy races a consumer already reading partial output; no commit barrier protects the read.
Push-based shuffle merge	incompatible	Push-based shuffle merges each partition's blocks into large pre-sealed files on merger services, and those files become readable only after a post-map-completion "finalize" step — the exact opposite of a pipelined shuffle, whose consumer reads before the producer finishes, so the two can never apply to the same edge.
Statically-indeterminate producer	moot	Its recovery is stage rollback-and-recompute, which a group never performs (§6: any failure aborts the group). The mechanism is never reached; rejected so the inapplicability is explicit rather than latent. (A producer whose output RDD is classified INDETERMINATE — a rerun can yield different data, not just a different order — determinable from the RDD graph at stage-creation time.)
Checksum-mismatch full retry	moot	The runtime counterpart to static indeterminism: it checksums each map task's output and, on a cross-attempt mismatch, rolls back and re-runs the succeeding stages. A group never keeps succeeding stages across a retry (§6), so this never fires; rejected to keep the inapplicability explicit.
Cached/persisted RDD in a member's within-stage chain	incompatible	Would capture partial output read mid-run and serve it as a complete result on reuse. Scoped to the member's within-stage (narrow-dependency) chain: a cached complete input reached via a broadcast variable or across a materialized shuffle — e.g. the static side of a stream-static join — is outside that chain and is unaffected.
Pipelined producer shared across jobs	incompatible	Breaks single ownership: a failure must map to exactly one job's group.
Regular shuffle internal to a group	incompatible	Concurrent co-scheduling contradicts materialize-before-read (§7).
Adaptive Query Execution over a pipelined exchange	incompatible	AQE reshapes exchanges from complete map-output statistics, which are unavailable while the shuffle is read incrementally. Enforced where exchanges are marked pipelined.

The AQE row forbids AQE's intra-batch, mid-stage replanning — reading a running producer's statistics
within one execution and reshaping its consumers — not statistics-driven adaptivity in general. A
streaming consumer may still reshape a later batch's plan from an earlier batch's completed statistics:
those statistics are final (the earlier batch fully materialized), and the reshaping happens at planning
time, before the later batch's stages are co-scheduled — so it composes with pipelining rather than
contradicting it. That cross-batch feedback is the compatible way to get AQE-like benefits (partition
coalescing, skew handling, join-strategy selection) in a pipelined streaming query.

[jerrypeng]

@cloud-fan @mridulm @viirya thank you for your responses. Based on your feedback, I have drafted a more formal spec of the behavior for the new scheduler primitive to be added:

#56055 (comment)

PTAL and let me know your thoughts. Thanks in advance!

[cloud-fan]

The spec is the right shape and answers the direction well — PSD as a binding contract (§2.1), the group as the failure/admission unit (§2.2), and activation by membership (§8, which removes the prototype's half-activated state). A few gaps before it's the implementation contract:

1. Cross-group scheduling / starvation (§4) — @mridulm's (b), and the weakest part. §4 defines gang admission for one group but not how multiple groups (or a group vs. regular jobs) arbitrate — and a pipelined group holds its slots for the whole batch. Is admission FIFO? Can a large group be starved by a stream of small ones? When an admitted group has pinned its slots, does the next one wait or fail-fast (which means a second streaming query can't start on a busy cluster)? Are running groups' slots subtracted from "available" for the next check?

2. "Available slots" needs a precise definition (§4). The prototype used defaultParallelism(), which returns spark.default.parallelism, not free slots — the barrier precedent uses sc.maxNumConcurrentTasks(rp) (checkBarrierStageWithNumSlots), which is the number to reuse. Also, sum(numTasks) ignores resource profiles: either state a group is single-resource-profile (and reject otherwise) or define per-profile accounting.

3. Output-commit vs. group-atomic failure (§5). §5 runs output-commit immediately while deferring stage/job finish, so a result stage → task commits → sibling fails in the replay window → group reruns → double commit. RTM's sink is idempotent so RTM is safe, but as a generic primitive this needs either output-commit to also defer, or a stated requirement that in-group result-stage side effects be idempotent.

4. Drop the fetch-failure framing in §6 — state it as a mechanism. Internal edges are streaming shuffle, so the base FetchFailed-means-resubmit concept doesn't apply inside a group; enumerating it implies a mechanism that doesn't exist. Better: any member task failure, for any reason, fails the group; the base single-stage resubmit path is disabled for members (that clause is what makes "all errors fail the group" actually true). Then note under §7 that a lost external durable input reruns the group rather than being recomputed in isolation — correct but coarser than base Spark, worth flagging as an intentional v1 simplification.

5. Cross-job stage reuse (§3/§4). Each micro-batch is a new job, but the base scheduler reuses shuffle-map stages via shuffleIdToMapStage. What forces a pipelined edge to create a fresh stage rather than binding batch N to batch N-1's cached stage?

Minor: (a) fail-fast admission → for streaming, crash-loops on restart; the barrier path retries BarrierJobSlotsNumberCheckFailed N times — say whether admission is retried or terminal. (b) "pipelined" over StreamingShuffleDependency is the better name (not streaming-specific) — worth a sentence on why.

[tgravescs]

Overall I really like the new idea of having a Pipelined Shuffle Dependency and support it directly in DagScheduler and define the contract. This opens it up to be generally useful in other scenarios. I'm still going through some of the details but here are a few initial comments

I think we need to clarify some of the gang scheduling wording.

A group is admitted only if the cluster can currently run all tasks of all member stages concurrently;

This is a group with multiple pipelined stages (pipelined group), not normal group with a single stage (ie how it runs today) Or is this specifically specified somewhere that we need to do Gang Scheduling?

I think this opens up some possibilities in the scheduler but I'm not sure it also addresses being able to have multiple shuffle managers or multiple ways to do the shuffle. If you are mixing groups where one is regular and the other is a pipelined group, within the pipelined group you would want to use the StreamingShuffleManager but in the regular group you would use the regular shuffle manager. I think you can make a MultiShuffleManager like is in the streaming shuffle proposal but I think within that it needs to be able to dynamically change per group. It seems like the ShuffleDependency should have a way to indicate what type of shuffle is required. This might be getting more in the future but thinking about it here if we are defining contracts and new ShuffleDependencies.

[jerrypeng]

@mridulm @tgravescs @cloud-fan @viirya

I created a PR that contains the spec to facilitate our discussion:

#57092

lets continue the discussion over there. I will close this PR.