viirya commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3817077723
##########
core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala:
##########
@@ -2604,6 +2687,60 @@ private[spark] class DAGScheduler(
.getOrElse(new Properties())
addPySparkConfigsToProperties(stage, properties)
+ // For a pipelined PRODUCER stage, tell its tasks which reduce partitions
the job actually
+ // reads (the result stage's partitions). The in-process channel writer
drops records
+ // routed to partitions no consumer will drain -- otherwise a partial-read
job (LIMIT /
+ // executeTake reads a subset) fills the unread partitions' bounded queues
and deadlocks
+ // the writer. The result stage is created before submitStage, so its
partitions are known
+ // here. Only the map-side producer needs this; a regular full-read job's
result stage
+ // covers every partition, so the property (if written) lists them all and
drops nothing.
+ // The live set is per-SHUFFLE-EDGE, not per-job: it is the reduce
partitions the
+ // consumer of THIS shuffle reads. It equals the job's result partitions
ONLY for the
+ // producer whose shuffle the result stage reads DIRECTLY (result
partition i maps to
+ // that producer's reduce partition i). A middle pipelined exchange in a
chain (e.g. a
+ // subquery's hash below a single-partition agg) is consumed by another
map stage that
+ // reads ALL its partitions, so it must stay fully live -- setting the
result's subset
+ // there would make it drop partitions the downstream stage still needs,
deadlocking.
+ // So set the property only on the result-feeding producer.
+ stage match {
+ case sms: ShuffleMapStage if isPipelinedProducer(stage) =>
+ // Notify the pipelined manager that this producer stage is being
(re)submitted, before
+ // any of its map tasks start. A transport with per-run state keyed by
shuffleId resets it
+ // here -- the one point with no live task of the new run, so the
reset cannot race the
+ // run's own writers/readers. The in-process channel transport clears
a prior run's
+ // abandoned-partition marks; the RPC streaming manager keeps no such
state (no-op default).
+ // The scheduler stays transport-agnostic: it calls the
PipelinedShuffleManager trait, not
+ // a concrete transport.
+ SparkEnv.get.pipelinedShuffleManager
+ .onPipelinedProducerStageSubmit(sms.shuffleDep.shuffleId)
+ jobIdToActiveJob.get(jobId).map(_.finalStage).collect { case rs:
ResultStage => rs }
+ .foreach { rs =>
+ // The live set is the result stage's partition ids. Those are the
shuffle's reduce
+ // partition ids -- so partition i means reduce partition i --
ONLY when the result RDD
+ // reaches this shuffle through an IDENTITY-PRESERVING chain:
every hop a 1:1,
+ // same-index OneToOneDependency (e.g. the mapPartitions wrapper
executeTake/collect
+ // adds over the ShuffledRowRDD, which preserves partition count
and index). It is NOT
+ // enough for the shuffle to be merely reachable: a narrow
operator that REMAPS
+ // partitions (coalesce -> a custom NarrowDependency, union -> a
RangeDependency offset)
+ // makes result partition ids differ from reduce partition ids,
and treating them as
+ // the live set would drop partitions a downstream operator still
pulls -- hanging the
+ // reader. When the chain is not identity-preserving the property
is left unset (all
+ // partitions live) and that operator drains every reduce
partition.
+ def readsShuffleByIdentity(rdd: RDD[_]): Boolean =
rdd.dependencies match {
Review Comment:
Good catch on the offset spec. `readsShuffleByIdentity` now also requires
`rdd.partitions.length == sd.partitioner.numPartitions` at the shuffle node, so
an offset `CoalescedPartitionSpec` (which passes the width-1 `require` but
breaks index identity) degrades to "all partitions live" instead of dropping a
partition and hanging. As you note nothing produces that today, so this is
defense-in-depth.
##########
core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala:
##########
@@ -85,7 +85,13 @@ private class DefaultRecordingManager(conf: SparkConf,
isDriver: Boolean)
override def shuffleBlockResolver: ShuffleBlockResolver =
mock(classOf[ShuffleBlockResolver])
}
private class IncrementalRecordingManager(conf: SparkConf, isDriver: Boolean)
- extends RecordingShuffleManager(conf, isDriver) with PipelinedShuffleManager
+ extends RecordingShuffleManager(conf, isDriver) with PipelinedShuffleManager
{
+ // A non-streaming pipelined manager (an in-process transport): it needs no
+ // StreamingShuffleOutputTracker, unlike the RPC streaming manager that the
trait defaults to.
+ // The "SparkEnv does not initialize the tracker when the incremental
manager is not streaming"
+ // and fail-loud tests rely on this manager reporting false here.
+ override def usesStreamingShuffleOutputTracker: Boolean = false
Review Comment:
Fixed. Added the positive case to `PipelinedShuffleRoutingSuite`: a
pipelined dep with a no-tracker manager routes to that manager, creates no
`StreamingShuffleOutputTracker`, and is absent from the `MapOutputTracker` --
pinning the relaxed successor to the deleted invariant, and exercising the
`None` arm of `outputTrackerMaster`.
--
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]
