dongjoon-hyun commented on PR #58097:
URL: https://github.com/apache/spark/pull/58097#issuecomment-5499453217
Thanks for the updates, @viirya. CI is green and I re-read the whole diff on
`b74f813`. Two new items (1, 2) are interactions with the existing Real-Time
Mode path that I missed in earlier rounds; item 3 is a regression in a fix from
the last round; the rest are smaller.
## 1. `EnablePipelinedShuffle` is applied to streaming plans too
`IncrementalExecution.preparations` is `state +: (super.preparations :+
PrepareTransformWithStateInitialStateForRealTimeMode :+
MarkPipelinedShuffleForRealTimeMode)`, and `super.preparations` is
`QueryExecution.preparations(...)`, which this PR now ends with
`Seq(EnablePipelinedShuffle())`. The rule has no streaming guard, so in local
mode with the flag and the channel manager on, every exchange in a micro-batch
or RTM plan -- the state-store shuffles, the static side of a stream-static
join -- is flipped to `pipelined = true` *before*
`MarkPipelinedShuffleForRealTimeMode` runs.
That contradicts what `markStreamingPath` deliberately does: it leaves the
static side regular because "marking it pipelined would pull it into the group
and demand slots for stages that must instead finish, which fails admission".
With this PR, the blanket rule has already marked it by the time RTM looks.
The rule (or `PipelinedShuffleEligibility`) needs to bail out for a
streaming `QueryExecution`, or `IncrementalExecution` needs to drop it from the
inherited list. Either way, a test that a streaming query in a feature-on
session is left untouched would pin it.
## 2. The new DAGScheduler paths run for the RPC streaming transport as well
Everything this PR adds to the scheduler is keyed on
`PipelinedShuffleDependency` / `isPipelinedProducer`, not on the manager, so it
also runs for Real-Time Mode jobs on the default streaming manager:
- `handleJobSubmitted` now clones the job `Properties` and stamps
`SPARK_PIPELINED_RUN_EPOCH` for every RTM job.
- `submitMissingTasks` runs `liveReduceSet` + `rddReachesShuffle` (and, on
`None`, a second `rddReachesShuffle`) for every RTM producer stage, and stamps
`SPARK_PIPELINED_LIVE_REDUCE_PARTITIONS` into its task properties.
- On an unmappable partial read it calls `abortStage(...)` with a message
telling the user to "disable spark.sql.shuffle.localPipelined.enabled" --
advice that does not apply to RTM.
The streaming writer never reads either property (nothing under
`core/.../shuffle/streaming` references them), so on that path this is pure
overhead plus a new failure mode, none of it exercised by the streaming suites.
The `requiresDetachedRecords` / `usesStreamingShuffleOutputTracker` pattern
already in this PR is the right shape: a capability on
`PipelinedShuffleManager` (say `supportsLiveReducePartitionHints`, default
`false`) that gates the live-set stamping and the abort, leaving the RTM
scheduler path byte-for-byte as it is on master.
## 3. The `classifyJobShuffleShape` early-out never fires
From the last round:
> `classifyJobShuffleShape` returns immediately when no pipelined manager is
configured, so a default cluster no longer pays for the graph walk itself on
every job submission.
```scala
// DAGScheduler.scala:1170
if (SparkEnv.get == null || SparkEnv.get.pipelinedShuffleManager == null) {
```
`spark.shuffle.manager.incremental` defaults to `"streaming"`, and
`SparkEnv.initializeShuffleManager()` instantiates it unconditionally --
`SparkEnv`'s own comment in `initializeStreamingShuffleOutputTracker` says "the
manager is non-null here". So on a driver this branch is unreachable and every
job submission on every deployment pays the new `(RDD, Boolean)`-keyed walk (a
tuple allocation per visit plus the boundary `HashMap`), which is strictly more
than the old `classifyJobShuffleKinds` over `traverseRDDGraph`'s
`HashSet[RDD[_]]`. The comment claiming otherwise should not ship either way.
A cheap fix that keeps the structure: run the old single-context walk first
(`hasPipelined` / `hasRegular` only), and run the two-context walk and the
materialization loop only when both are true. All-regular and all-pipelined
jobs then cost exactly what they cost on master.
## 4. `ContextCleaner` names a concrete transport
```scala
// ContextCleaner.scala:306
val mgr = SparkEnv.get.pipelinedShuffleManager
mgr != null && !mgr.usesStreamingShuffleOutputTracker &&
org.apache.spark.shuffle.local.pipelined.ChannelShuffleRendezvous.holdsShuffle(shuffleId)
```
Removing the `holdsShuffle` trait method fixed the leak but made core's
generic cleanup depend on one transport by fully-qualified name. The registry
was the bug, not the abstraction: `def holdsShuffle(shuffleId: Int): Boolean =
false` on `PipelinedShuffleManager`, with the channel manager delegating to the
rendezvous, gives the same behavior and keeps `ContextCleaner`
transport-agnostic. Also `ContextCleaner` already holds `sc` and uses
`sc.env.broadcastManager` below; `sc.env.pipelinedShuffleManager` is the
consistent form (same for `DAGScheduler.outputTrackerMaster`, which mixes
`SparkEnv.get` and `sc.env` on adjacent lines).
## 5. The rendezvous key forces a full scan per cleanup
`queues` and `abandoned` are keyed by the flat `(shuffleId, epoch,
reducePartitionId)`, so `holdsShuffle` and `removeShuffle` iterate every live
entry. `holdsShuffle` is now called from `doCleanupShuffle` for **every**
shuffle cleaned in a feature-on session, including the regular prefix shuffles
this feature produces, so a long-lived local session pays O(N x live entries).
Keying by `shuffleId` first makes both O(1) and makes `removeShuffle` one
atomic `remove`.
## 6. `liveReduceSet`'s identity check is a count heuristic, and a wrong
answer is silent data loss
```scala
if (cur.partitions.length == sd.partitioner.numPartitions) resultReduce ++=
live
else unmappable = true
```
Equal counts do not imply an identity index mapping (one skew-split plus one
coalesce of two gives the same length and a different mapping). Today the AQE
rule keeps pipelined exchanges out of `ShuffleQueryStageExec` so neither
`OptimizeSkewedJoin` nor `CoalesceShufflePartitions` can produce that -- but
that is a scheduler-level invariant resting on a SQL-rule behavior, with no
check between them. What makes it worth tightening is the failure mode: the
writer's `liveMask` silently drops every record routed outside the computed
set, so an under-approximation yields wrong results, not the hang the mechanism
exists to prevent. Everything else in this PR fails loud; this path fails
quiet. At minimum, count dropped records in the writer and log them at debug;
better, make the identity property explicit (restrict the match to
`ShuffledRDD` / width-1 `CoalescedPartitionSpec`, or have the reader RDD report
the reduce id for a partition index).
## 7. Reader fetch-wait time truncates to zero on fast polls
```scala
// ChannelShuffleWriterReader.scala:292
readMetrics.incFetchWaitTime((System.nanoTime() - start) / 1000000L)
```
This converts each `poll` interval to ms separately; a normal hand-off
returns in microseconds, so every sub-millisecond wait contributes 0 and a
consumer that is genuinely waiting reports ~0 -- the opposite of the intent in
the comment above it. Accumulate nanos and convert once, as
`ShuffleBlockFetcherIterator.withFetchWaitTimeTracked` does. (The write side is
fine: `writeTime` is in nanoseconds.)
## 8. The benchmark cannot be regenerated by the project's benchmark workflow
`runBenchmarkSuite` returns immediately when `cores < 11`.
`.github/workflows/benchmark.yml` runs on `ubuntu-latest` (4 cores), so the
workflow would print `[skip]` and never write a results file -- which is why
the checked-in file was produced on an "Apple M4 Max / Mac OS X 26.5.2" rather
than by the workflow, and why only `-jdk21-results.txt` exists where every
other benchmark under `sql/core/benchmarks` has `-results.txt`, `-jdk21-`, and
`-jdk25-`. Either shrink the shapes so the gang fits the standard runner (it is
`inputParts + 8 + 1`; with `spark.sql.shuffle.partitions=2` the minimum shape
fits in 5 slots), or don't check in results that no committer can reproduce
through the normal path.
## 9. Config naming and surface
- The streaming transport's knobs are `spark.shuffle.streaming.*`; for
symmetry the channel transport's should be `spark.shuffle.channel.batchSize` /
`spark.shuffle.channel.queueCapacity`, not `spark.shuffle.pipelined.channel.*`.
- `spark.sql.shuffle.localPipelined.enabled` is public, documented
"Experimental.", and does nothing unless `spark.shuffle.manager.incremental` is
also set to the channel manager -- the mismatch is only a `logDebug`. Either
mark it `.internal()` while the feature bakes, or make the flag self-sufficient
and warn once when the manager does not match.
## 10. Small things
- `EnablePipelinedShuffle`'s scaladoc says "a production version would be a
targeted, cost/shape-aware replacement rather than a blanket rewrite". Either
the blanket rule is what we ship (say so), or that sentence names a follow-up
JIRA; it should not ship as-is.
- `EnablePipelinedShuffle()` and `AQEEnablePipelinedShuffle()` are no-field
case classes; `object` is the usual form.
- `ShuffleExchangeExec.stringArgs`'s comment, "`pipelined` is only
meaningful for a Real-Time Mode plan", is no longer true after this PR.
- `PipelinedShuffleSqlSuite`, `AQEPipelinedShuffleSuite`, and
`PipelinedLimitHangSuite` each carry a near-identical ~30-line session harness,
and each begins by stopping and clearing whatever active/default `SparkSession`
the JVM holds. I could not find precedent for that under `sql/core/src/test`;
please factor it into one trait and reconsider tearing down a session the suite
did not create.
- The new files run 46-65% comment lines, and several narrate review history
("This replaces an earlier design that...", "An early version recorded each
shuffle's map-task count..."). That belongs in the JIRA and commit message; the
genuinely subtle invariants (epoch keying, width-1 reads, the abandon/offer
re-check) are worth every line.
--
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]