jerrypeng commented on code in PR #57092:
URL: https://github.com/apache/spark/pull/57092#discussion_r3584925481
##########
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
Review Comment:
Sounds good. Let me add this information explicitly
--
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]