jerrypeng commented on code in PR #57092:
URL: https://github.com/apache/spark/pull/57092#discussion_r3584997426


##########
PIPELINED_SHUFFLE_DEPENDENCY_SPEC.md:
##########
@@ -0,0 +1,409 @@
+# 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, when two stages of a job are connected by a shuffle, they run one after 
another: the consumer
+stage does not start until the producer's shuffle output is fully 
materialized. (Independent stages —
+ones not connected by a shuffle — already run concurrently.) Some workloads 
need even
+shuffle-connected stages to run **concurrently**, with the consumer reading 
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. 
("Pipelined",
+  "incremental", and "transient" all describe this same edge in this doc — 
read as its output being
+  streamed to the consumer as produced, not stored.)
+- 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.
+- It is also the **per-dependency selector** for the shuffle implementation: 
the shuffle layer maps a
+  pipelined dependency to an incremental `ShuffleManager` and everything else 
to the default, so one
+  job with both regular and pipelined groups uses the right implementation for 
each. The scheduler
+  construct stays generic; the shuffle implementation stays pluggable.
+  - For example, an additional conf (e.g. `spark.shuffle.manager.incremental`) 
can be introduced to
+    specify the incremental shuffle implementation used for pipelined shuffle 
dependencies, alongside
+    the existing `spark.shuffle.manager` for the default.
+
+**What the subtype inherits vs. what the scheduler skips.** A pipelined 
shuffle dependency is modeled
+as a `ShuffleDependency` subtype, which cleanly divides `ShuffleDependency`'s 
behavior in two — an
+invariant every component that handles the type must preserve:
+- *Inherited — it is treated as a real shuffle dependency:* (1) 
**stage-boundary splitting** —
+  `getShuffleDependenciesAndResourceProfiles` matches `case shuffleDep: 
ShuffleDependency`
+  (`DAGScheduler.scala:882`), so the subtype splits producer and consumer into 
separate stages
+  automatically; this is what §3 ("introduces a shuffle boundary exactly as a 
regular one does")
+  relies on. (2) **transport** — `shuffleId`, `ShuffleWriter`/`ShuffleReader`, 
and `ShuffleManager`
+  registration, all from the base constructor.
+- *Scheduler-skipped — it must not be treated as a materialized shuffle for 
the post-boundary
+  lifecycle:* `MapOutputTracker` registration (§6), `shuffleIdToMapStage` 
reuse (§4), and
+  executor-loss output-removal-then-resubmit (§6). This is clean rather than 
"inherit then disable":
+  all three are `DAGScheduler`-driven at `createShuffleMapStage` 
(`DAGScheduler.scala:667-673`) and
+  gated per dependency, so a pipelined dependency is simply never opted into 
them.
+- *Match-site invariant:* a marker subtype relies on every existing `case _: 
ShuffleDependency` site
+  meaning "materialized" for the new type. The reuse/tracking/loss sites are 
safe because they are
+  scheduler-gated as above; the remaining sites were audited and treat a 
pipelined dependency as an
+  ordinary materialized shuffle wherever the subtype does not specifically 
diverge. Preserving this
+  is an invariant for implementers, not something to rediscover.
+
+A **regular shuffle dependency (RSD)** is an ordinary shuffle dependency: its 
output must be fully
+materialized before any consumer reads it.
+
+Note: the name *pipelined* is deliberately chosen over *streaming*. The 
property is that a consumer
+reads producer output as it is produced — software-pipelining of dependent 
stages — a general
+execution capability. Streaming / real-time mode (RTM) is the first caller, 
but nothing about the
+primitive is streaming-specific.
+
+### 2.2 Pipelined group (PG)
+
+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 PG:** a regular shuffle dependency whose consumer is in 
the group and whose
+producer is not — i.e. a normal materialized parent of the group.
+
+**Outputs of PG:** a `ResultStage` **may** be a member of a pipelined group — 
a pipelined consumer
+that produces the job's result (its commit handling is §5) — but a PG **need 
not** contain one.
+When it does not, its outputs to the rest of the DAG are regular 
(materialized) shuffle edges from
+its members to consumers outside the group, and there may be **more than one** 
— e.g. two members
+each feeding a distinct downstream stage. Such a PG is an interior component 
whose materialized
+outputs feed downstream groups (§7); in that respect it behaves like a barrier 
stage embedded in a
+larger DAG — an interior unit with regular-shuffle edges in and out.
+
+**Example.** Take a job with stage DAG `Z -> A -> B -> C`, where `Z -> A` and 
`B -> C` are regular
+edges and `A -> B` is pipelined. Grouping over pipelined edges gives one 
non-singleton group
+`PG = {A, B}` (`Z` and `C` are singletons). `Z`'s output is the group's 
external input, which the
+group waits for; `A` and `B` run concurrently, `B` reading `A`'s output as `A` 
produces it; and
+`B`'s output to `C` is a materialized output of the group that `C` consumes by 
normal
+materialize-before-read sequencing.
+
+---
+
+## 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.
+- **A DAG may contain multiple pipelined groups.** Disjoint sets of pipelined 
edges form distinct
+  groups, and a group's materialized output may feed another group (§7). 
Independent groups (no
+  dependency path between them) may be admitted and run concurrently, subject 
to the cross-group
+  slot arbitration in §4.1; a group that consumes another group's materialized 
output waits for it,
+  by the normal materialize-before-read sequencing on that regular edge. This 
mirrors how a DAG may
+  contain multiple barrier stages.
+
+---
+
+## 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 (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 pipelined 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. A pipelined group is never left with some members running while others 
wait on slots the
+  running members occupy.
+  - *A non-pipelined (singleton) group is unaffected:* it is admitted exactly 
as a normal stage is
+    today — submitted when its parents are available, filling slots 
incrementally, with no all-at-once
+    requirement. Gang admission applies only to a group of two or more stages 
connected by pipelined
+    edges.
+- **Slot check.** The group's aggregate concurrent-task demand — the sum of 
`numTasks` over member
+  stages — is compared against the number of **currently-free** slots: the 
cluster's total
+  concurrent-task capacity minus the tasks already running. If the group needs 
more than are free,
+  admission fails fast, since the group could never become fully co-resident. 
(Free rather than
+  total is essential once more than one group can be admitted; §4.1 explains 
why.)
+  - *How capacity is measured:* the total concurrent-task capacity is the 
value barrier's slot check
+    uses (`sc.maxNumConcurrentTasks` — for the group's resource profile, the 
number of task slots
+    across active executors, i.e. cores divided by cores-per-task), not 
`spark.default.parallelism`
+    (a default partition count, unrelated to how many tasks can run 
concurrently).
+- **Single resource profile per group (v1).** A resource profile is the 
executor/task resource
+  requirement (cores, memory, GPUs, ...) a stage runs under; the number of 
concurrent slots is
+  defined *per profile* (a cluster may run many concurrent CPU tasks but few 
GPU tasks). Comparing
+  one demand against one capacity is therefore only well-defined when all 
member stages of a group
+  share a single profile. v1 requires that and rejects a mixed-profile group 
(fail-fast, §9);
+  per-profile accounting — checking each profile's demand against that 
profile's own capacity — is a
+  follow-up, not needed for the streaming shapes whose members share the 
default profile.
+- **Co-residency.** Once admitted, all member stages of a group are 
simultaneously running.
+- **Single ownership and 1:1 (v1).** In v1 a pipelined producer feeds exactly 
one consumer and
+  belongs to exactly one job. These are two separate restrictions. They follow 
from what a pipelined
+  shuffle *is* (a transient edge), not from how any particular caller builds 
its DAG — the pipelined
+  dependency and PG are generic `DAGScheduler` constructs that code can depend 
on directly, not only
+  through SQL / real-time mode:
+  - *No cross-job / cross-time reuse — intrinsic.* A pipelined shuffle is 
transient: a once-through
+    live stream with no retained, addressable output. So, unlike a regular 
shuffle, there is nothing
+    for a second job to reuse — reuse across jobs is unsound for it.
+    - This must be enforced, not assumed: from the scheduler's perspective a 
shuffle-map stage *can*
+      be reused unless something forbids it. Spark would otherwise reuse a 
shuffle in two ways, and
+      **both** must be blocked. (1) Stage-object reuse via 
`shuffleIdToMapStage`: a pipelined
+      dependency is not enrolled in it, so each gets a fresh producer stage. 
(2) Output-availability
+      reuse: a re-created producer would be skipped as already-done because 
its outputs are still
+      registered in
+      `MapOutputTracker` (`getMissingParentStages` treats a registered stage 
as available). Blocking
+      only (1) is not enough. The fix for (2) is the same one §6 requires for 
executor loss: a
+      pipelined shuffle is not tracked in `MapOutputTracker` at all — its 
availability is owned
+      elsewhere — so neither cross-job reuse nor executor-loss removal can act 
on it.
+    - A pipelined producer whose stage carries more than one jobId is rejected 
(fail-fast, §9).
+    - (A producer may separately write a durable *materialized* output edge — 
a distinct regular
+      `ShuffleDependency` that reuses normally. Run-once binds only the 
transient edge.)
+  - *1:1 within the group (v1) — deferred, not intrinsic.* v1 rejects a 
pipelined producer with more
+    than one consuming stage, but 1:N fan-out is a supported model that v1 
defers rather than an
+    incompatibility (§9): co-scheduled consumers would read the live stream 
via multicast, and
+    non-co-scheduled consumers read a materialized edge the producer also 
writes.
+
+### 4.1 Cross-group admission (multiple groups / groups vs. regular jobs)
+
+A pipelined group holds its slots for its whole run, so admission is decided 
against currently-free
+slots.
+
+- **Why free slots, not total** (the slot check itself is defined in §4). Free 
capacity subtracts
+  the tasks already running — for every running group and regular job — so a 
group is admitted only
+  if its full demand fits in what the others leave free. This is what makes 
co-residency safe:
+  comparing against *total* capacity would let two groups each pass the check 
yet fail to co-fit,
+  which is exactly the partial co-residency that gang admission forbids. This 
is where the check
+  diverges from barrier, which compares demand against total capacity; 
computing free capacity is a
+  new computation, not barrier's check reused verbatim.
+- **No waiting queue, no partial reservation.** A group that doesn't fit fails 
its admission; it never
+  sits in a queue holding slots incrementally. This keeps the scheduler change 
minimal and cannot
+  deadlock (a group never occupies slots a sibling is blocked on).
+- **Two failure reasons, one path.** Demand > total capacity can never fit (a 
plan/sizing error);
+  demand > free slots but ≤ total could fit later. Both fail admission.
+- **Retry (v1: caller's decision; scheduler-side retry is a post-v1 
refinement).** In v1 the
+  scheduler does not retry a failed admission — it fails the job, and retry is 
the caller's decision
+  (for a streaming query, the batch-execution loop restarting the batch with 
its own backoff). This
+  is adequate only for a simple `prefix* -> PG -> suffix*` shape run by a 
caller that has its own
+  restart loop. Its unit is the whole job, so it has two weaknesses: with 
concurrent work (sibling
+  stages, other PGs) or an already-committed prefix, restarting the job 
discards work it cannot
+  selectively preserve; and a directly-depended-on PG may have no retrying 
caller at all. The result
+  is that a transient slot shortfall either kills the job or forces 
over-provisioning.
+  - *(Refinement, post-v1.)* The scheduler retries **admission of the PG 
specifically** — hold the
+    group, re-run the gang slot-check on a timer, admit when it fits, fail 
after a bounded number of
+    attempts — leaving any completed prefix and concurrent stages untouched. 
This mirrors barrier's
+    transient-shortfall retry (re-post on a scheduled interval bounded by a 
max-failure count, then
+    fail; `DAGScheduler` `BarrierJobSlotsNumberCheckFailed` path), the 
difference being that a
+    mid-DAG PG retries its own admission rather than re-posting the whole job 
(barrier can re-post
+    the job only because its check runs before any stage executes). Making 
transient-shortfall
+    tolerance a property of the construct, rather than a burden each caller 
re-implements, is why it
+    belongs in the scheduler.
+
+---
+
+## 5. Completion
+
+- **Group-observable completion.** A member stage is not observably finished 
until **all** member
+  stages of its group have completed successfully. A member that finishes 
ahead of the others cannot
+  advance stage or job completion, nor make its output visible to a downstream 
consumer, until the
+  group does — whether that output is a `ResultStage`'s job result or a 
materialized shuffle feeding a
+  downstream group (§2.2).
+- **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.
+  - *Deferred output registration.* A member's materialized output edge (to a 
consumer outside the
+    group, §7) is written as its tasks run, but its map-output registration is 
**deferred to group
+    completion** — so an out-of-group consumer, which waits for the group 
anyway (§7), only ever sees
+    the single committed version. A failed intermediate attempt's blocks are 
simply orphaned (never
+    registered), exactly as for a resubmitted regular map stage today. This is 
what makes a producer
+    that writes both an incremental (in-group) and a materialized 
(out-of-group) output edge
+    consistency-clean: the group's internal replay never reaches the 
materialized consumer, because
+    it has not started.
+- **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.
+- **In-group result-stage side effects must be idempotent.** Per-task side 
effects run immediately
+  (above), including a result stage's output commit. If a result task commits 
and a sibling then
+  fails in the replay window, the group reruns and re-delivers that output — 
the standard streaming
+  model, where a batch is re-delivered on recovery and the sink must absorb 
it. So v1 requires an
+  in-group result stage's side effects to be idempotent, and fail-fast rejects 
a sink that cannot
+  absorb re-delivery (§9).
+- **Group-atomic rerun must reset per-partition commit authorization.** A PG 
is an atomic
+  scheduling unit: there are no per-task or per-partition replays within it. 
`OutputCommitCoordinator`
+  is built on the opposite assumption — it authorizes exactly one attempt per 
partition and
+  permanently denies any later request for a partition that already committed. 
That is sound today
+  because a committed partition is never recomputed: a stage rerun (e.g. on 
fetch failure) recomputes
+  only its *missing* partitions, so a committed task never needs to commit 
again, and the permanent
+  deny is correct. A PG breaks that assumption — because the group is atomic, 
a failure anywhere
+  reruns the *whole* group, including a result stage whose tasks already 
succeeded and committed. Those
+  committed partitions are rerun and must be allowed to commit again, but the 
coordinator still holds
+  the previous attempt's authorized committers and would deny them (a 
`Success` clears nothing; only a
+  *failed* holder's slot is cleared today). So the group rerun must let the 
new tasks commit, by one
+  of: (a) rerunning members under **fresh stage ids** — a fresh coordinator 
`StageState` with no prior
+  committers; or (b) **resetting** the committed state for those stages in the 
`OutputCommitCoordinator`
+  (e.g. `stageEnd` on teardown) before the rerun.
+
+### 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 
group-atomic completion (the
+point at which the group's outputs become observable — distinct from the 
per-task output-commit
+above): **task-level events flow in real time, but stage-completion and 
job-completion events are
+held until the group completes — 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 completion; on group 
failure, emitted with a failure reason | "Completed" should track group 
completion, which is atomic at the group level. A member whose tasks finish 
early is reported as still running until the group completes — 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 completion only | Job 
completion delivers results to the caller and cancels sibling stages; emitting 
it before the group completes 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: 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 member task failure, for any reason, fails 
the whole group.
+  - *Single-stage resubmit is not taken for group members.* The 
isolated-resubmit branch is never
+    taken for a member of a pipelined group: the transient pipelined shuffle 
cannot be re-read and
+    members are co-scheduled, so a lone-stage resubmit is never valid — every 
member failure instead
+    routes to group failure.
+    - *Note this differs from the base scheduler,* which handles a task 
failure by resubmitting just
+      that one stage — the `FetchFailed` -> resubmit path, and retry paths 
generally — recomputing
+      the failed stage in isolation and resuming.
+  - *A pipelined shuffle's availability is not tracked in `MapOutputTracker`.* 
This is the mechanism
+    that makes "no single-stage resubmit" robust, and it also closes an 
executor-loss resubmit path.
+    A regular shuffle records its map outputs in `MapOutputTracker`, and 
`ShuffleMapStage.isAvailable`
+    is derived from them (`getNumAvailableOutputs`); on executor or host loss 
the scheduler strips
+    those outputs (`removeOutputsOnExecutor` / `removeOutputsOnHost`), which 
flips `isAvailable` to
+    false and resubmits the producer. For a transient pipelined shuffle, 
resubmitting the producer is
+    both wrong (its output was never durably stored, so there is nothing to 
recompute into) and
+    dangerous (the resubmitted producer has no live consumers left to serve 
and can hang — for
+    example, an incremental-shuffle writer that blocks waiting for 
end-of-stream acknowledgements from
+    reducers that have already finished). So a pipelined shuffle is *not* 
registered with
+    `MapOutputTracker`; its map-stage availability is
+    tracked separately — on the `ShuffleMapStage` itself (a monotonic 
completed-partition set) or a
+    streaming-specific tracker (`StreamingShuffleOutputTracker`) — and 
`MapOutputTracker` stays
+    streaming-agnostic. Executor/host-loss removal then never touches a 
pipelined shuffle and never
+    flips its availability, so the producer is never resubmitted from that 
path; genuine losses that
+    happen while the group is still running are handled by group-atomic 
teardown above. This is also
+    what makes the "no single-stage resubmit" rule above robust: with no 
tracked availability to flip,
+    the availability-driven resubmit cannot fire in the first place. (Same
+    `MapOutputTracker`-must-not-govern-a-transient-shuffle point as the reuse 
layer in §4, seen from
+    the executor-loss side rather than the cross-job-reuse side.)
+  - *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 it is normal, not rare, for the producer to be 
available while the
+    consumer is still running. 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.
+  - *(Refinement, post-v1.)* v1 fails the group on any member's executor loss 
— safe and hang-free,
+    but over-eager: a producer that finished and whose output was already 
fully consumed can lose its
+    executor with no impact on the DAG, yet still triggers a rerun. A more 
precise, consumer-driven
+    signal is a post-v1 optimization: have the streaming reader raise a 
`FetchFailed` only when it
+    actually fails to read its input, so a post-consumption producer loss 
yields no failure and no
+    teardown. This reuses Spark's existing `FetchFailed` channel as the 
notification, with one
+    adaptation — on a pipelined edge it must route to **group failure**, not 
the base single-stage
+    mapper-resubmit — so detection is consumer-driven while recovery stays 
group-atomic.
+- **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.
+  - A lost external durable input reruns the group rather than being 
recomputed in isolation.
+- **Regular shuffle out of a group — supported.** A member may have a regular 
output edge to a stage

Review Comment:
   Agreed. Though in theory the pipelined group should just rerun as a whole 
given its inputs are still available just as a normal stage.
   
   Let me add to clarify



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to