dongjoon-hyun opened a new pull request, #59738:
URL: https://github.com/apache/spark/pull/59738
### What changes were proposed in this pull request?
This PR makes `DAGScheduler.handleMapStageSubmitted` reject a map-stage job
whose shuffle dependency has a `PipelinedShuffleDependency` anywhere in its RDD
lineage. The job fails with `PIPELINED_SHUFFLE_UNSUPPORTED` before any stage is
created, in the same way that `handleJobSubmitted` rejects the same graph for a
result job.
```scala
if (rddGraphHasPipelinedDependency(dependency.rdd)) {
val e = pipelinedUnsupportedError("a map-stage job that depends on a
pipelined shuffle")
logWarning(...)
listener.jobFailed(e)
return
}
```
This adds one RDD graph walk per map-stage job, similar to the walk that
`handleJobSubmitted` already does for every result job.
### Why are the changes needed?
`handleMapStageSubmitted` rejected only a dependency that is itself a
`PipelinedShuffleDependency`. A regular `ShuffleDependency` whose map side
reads a pipelined shuffle, directly or further down its lineage, was accepted.
Such a job skipped every pipelined job check of `handleJobSubmitted`:
- the shuffle shape check, which rejects a pipelined shuffle below a regular
shuffle
- gang admission (`rejectUnadmittablePipelinedGroup`)
- `checkPipelinedJobSupported` and the speculation check
- the run epoch (`SPARK_PIPELINED_RUN_EPOCH` and
`PipelinedShuffleManager.startRun`)
- `ActiveJob.hasPipelinedDependency`
However, `submitStage` still co-scheduled the pipelined producer with the
stage reading it, as ordinary task sets (`TaskSet.isPipelined = false`).
With the in-process channel transport (`PipelinedChannelShuffleManager`),
this causes a hang or wrong results.
- **Hang.** The tasks have no run epoch, so every channel access throws
`TaskKilledException("Pipelined shuffle run has ended")`. A kill does not count
as a task failure in a non-pipelined task set, so the tasks are relaunched
forever. In a `local[4]` test, 73,470 task attempts were killed in 10 seconds,
each with a WARN log line, and the job never finished.
- **Correctness.** If run epoch 0 is active, because job 0 of the
`SparkContext` is a pipelined job that is still running, the channel accesses
succeed. A producer attempt that is retried, or relaunched after a kill, then
delivers its rows again. The channel reader counts end-of-stream markers
without tracking writer attempts, so the map-stage job can succeed with
duplicated rows, or stop early and drop the rows of other producers. In local
tests, the reader saw 5,120 rows instead of 3,072 when a producer attempt
failed after two batches. When a producer attempt failed after its
end-of-stream marker, the reader counted 200 rows from that producer and none
from the other, instead of 100 each. Both map-stage jobs reported success.
With the RPC streaming transport, the job runs without gang admission or
group-atomic failure handling, and speculation is not rejected.
`SparkContext.submitMapStage` is `private[spark]`, and its only in-tree
caller is AQE (`ShuffleExchangeExec`). `AQEEnablePipelinedShuffle` keeps each
plan all-pipelined or all-regular, so SQL does not create this shape within one
plan. It does when an RDD that contains a pipelined shuffle is wrapped in a new
DataFrame and shuffled under AQE. For example, in local mode with
`spark.shuffle.manager.incremental=org.apache.spark.shuffle.local.pipelined.PipelinedChannelShuffleManager`
and `spark.sql.shuffle.localPipelined.enabled=true`:
```scala
val inner = spark.range(0, 1000, 1, 2).withColumn("k", $"id" %
7).groupBy($"k").count()
val rows = inner.queryExecution.toRdd.map(r => Row(r.getLong(0),
r.getLong(1)))
spark.createDataFrame(rows, inner.schema).groupBy($"count").count().collect()
```
A user RDD that creates a `PipelinedShuffleDependency` (a `DeveloperApi`)
reaches it the same way.
### Does this PR introduce _any_ user-facing change?
Yes, within the unreleased branches only. A map-stage job that depends on a
pipelined shuffle now fails up front instead of hanging or completing with
wrong results. For example, the query above fails with:
```
[PIPELINED_SHUFFLE_UNSUPPORTED] A pipelined shuffle stage group cannot be
scheduled because it uses an unsupported feature: a map-stage job that depends
on a pipelined shuffle. ...
```
No released Spark version has `PipelinedShuffleDependency`. `branch-4.3` has
it with the RPC streaming transport, and `branch-4.x` and `master` also have
the in-process channel transport.
### How was this patch tested?
Pass the CIs with the newly added test cases.
- `DAGSchedulerSuite`: the job is rejected when its map stage reads a
pipelined shuffle directly and when the pipelined shuffle sits below another
regular shuffle, and no task set or partial state is left.
- `PipelinedChannelShuffleSuite`: the job is rejected before any task
starts, and the same pipelined shuffle still runs as a result job.
- `AQEPipelinedShuffleSuite`: the query above fails fast instead of hanging,
and the inner pipelined query still runs.
Without the fix, the first test fails because the job is accepted, and the
other two time out. The duplicated and dropped rows above were reproduced with
temporary tests that keep job 0 running, which are not included.
```
$ build/sbt 'core/testOnly *DAGSchedulerSuite *PipelinedChannelShuffleSuite
-- -z SPARK-60497'
$ build/sbt 'sql/testOnly *AQEPipelinedShuffleSuite -- -z SPARK-60497'
```
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 5.5
--
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]
